Merge pull request #97 from deel-ai/feat/test_uniformization

Merge the unit tests using pytest. All tests are common with keras3 branch, and also torchlip library (https://github.com/deel-ai/deel-torchlip)
The only difference is in the file:
tests/utils_framework.py

    Modify Readme.md + docs/index.md to change the banners
    increment to minor version 1.5.1

Author: cofri

.github/workflows/python-tests.yml

@@ -17,7 +17,7 @@ jobs:
       matrix:
         include:
           - python-version: 3.8
-            tf-version: 2.3
+            tf-version: 2.4
           - python-version: 3.9
             tf-version: 2.7
           - python-version: "3.10"

CONTRIBUTING.md

@@ -32,7 +32,7 @@ $ make test
 
 This command will:
 - check your code with black PEP-8 formatter and flake8 linter.
-- run `unittest` on the `tests/` folder with different Python and TensorFlow versions.
+- run `pytest` on the `tests/` folder with different Python and TensorFlow versions.
 
 
 ## Submitting your changes

README.md

@@ -1,8 +1,8 @@
 <div align="center">
         <picture>
-                <source media="(prefers-color-scheme: dark)" srcset="./docs/assets/logo_white.svg">
-                <source media="(prefers-color-scheme: light)" srcset="./docs/assets/logo.svg">
-                <img alt="Library Banner" src="./docs/assets/logo.svg">
+                <source media="(prefers-color-scheme: dark)" srcset="./docs/assets/banner_dark_deellip.png">
+                <source media="(prefers-color-scheme: light)" srcset="./docs/assets/banner_light_deellip.png">
+                <img alt="DEEL-LIP Banner" src="./docs/assets/banner_light_deellip.png">
         </picture>
 </div>
 <br>
@@ -44,9 +44,7 @@ layers for keras**.
 > **Incompatibility with TensorFlow >= 2.16 and Keras 3**
 >
 > Due to significant changes introduced in TensorFlow version 2.16 and Keras 3, this
-> package is currently incompatible with TensorFlow versions 2.16 and above. Users are
-> advised to use TensorFlow versions lower than 2.16 to ensure compatibility and proper
-> functionality of this package.
+> 'master' branch is currently incompatible with TensorFlow versions 2.16 and above. 
 >
 > We're excited to announce that a nightly version of deel-lip with support for Keras 3
 > is now available for testing! If you'd like to experiment with the latest updates
@@ -136,7 +134,8 @@ More from the DEEL project:
 
 - [Xplique](https://github.com/deel-ai/xplique) a Python library exclusively dedicated to explaining neural networks.
 - [Influenciae](https://github.com/deel-ai/influenciae) Python toolkit dedicated to computing influence values for the discovery of potentially problematic samples in a dataset.
-- [deel-torchlip](https://github.com/deel-ai/deel-torchlip) a Python library for training k-Lipschitz neural networks on PyTorch.
+- [deel-TorchLip](https://github.com/deel-ai/deel-torchlip) a Python library for training k-Lipschitz neural networks on PyTorch.
+- [Oodeel](https://github.com/deel-ai/oodeel) a Python library for post-hoc deep OOD (Out-of-Distribution) detection on already trained neural network image classifiers
 - [DEEL White paper](https://arxiv.org/abs/2103.10529) a summary of the DEEL team on the challenges of certifiable AI and the role of data quality, representativity and explainability for this purpose.
 
 ## 🙏 Acknowledgments

deel/lip/VERSION

@@ -1 +1 @@
-1.5.0
+1.5.1

deel/lip/layers/pooling.py

@@ -22,6 +22,7 @@
 """
 
 import numpy as np
+from typing import Tuple
 import tensorflow as tf
 import tensorflow.keras.layers as keraslayers
 from tensorflow.keras.utils import register_keras_serializable
@@ -389,39 +390,51 @@ def __init__(
             The image shape must be divisible by the pool shape.
 
         Args:
-            pool_size: tuple describing the pool shape
+            pool_size: int or tuple describing the pool shape (p_h, p_w)
             data_format: can either be `channels_last` or `channels_first`
             name: name of the layer
             dtype: dtype of the layer
             **kwargs: params passed to the Layers constructor
         """
         super(InvertibleDownSampling, self).__init__(name=name, dtype=dtype, **kwargs)
-        self.pool_size = pool_size
         self.data_format = data_format
 
-    def call(self, inputs):
-        if self.data_format == "channels_last":
-            return tf.concat(
-                [
-                    inputs[
-                        :, i :: self.pool_size[0], j :: self.pool_size[1], :
-                    ]  # for now we handle only channels last
-                    for i in range(self.pool_size[0])
-                    for j in range(self.pool_size[1])
-                ],
-                axis=-1,
-            )
+        ndims = 2
+        ks: Tuple[int, ...]
+        if isinstance(pool_size, int):
+            ks = (pool_size,) * ndims
         else:
-            return tf.concat(
-                [
-                    inputs[
-                        :, :, i :: self.pool_size[0], j :: self.pool_size[1]
-                    ]  # for now we handle only channels last
-                    for i in range(self.pool_size[0])
-                    for j in range(self.pool_size[1])
-                ],
-                axis=1,
+            ks = tuple(pool_size)
+
+        if len(ks) != ndims:
+            raise ValueError(
+                f"Expected {ndims}-dimensional pool_size, but "
+                f"got {len(ks)}-dimensional instead"
+            )
+        self.pool_size = ks
+
+    def call(self, inputs):
+        if self.data_format == "channels_first":
+            # convert to channels_first
+            inputs = tf.transpose(inputs, [0, 2, 3, 1])
+        # from shape (bs, w*pw, h*ph, c) to (bs, w, h, c, pw, ph)
+        input_shape = tf.shape(inputs)
+        w, h, c_in = input_shape[1], input_shape[2], input_shape[3]
+        pw, ph = self.pool_size
+        wo = w // pw
+        ho = h // ph
+        inputs = tf.reshape(inputs, (-1, wo, pw, h, c_in))
+        inputs = tf.reshape(inputs, (-1, wo, pw, ho, ph, c_in))
+        inputs = tf.transpose(
+            inputs, [0, 1, 3, 5, 2, 4]
+        )  # (bs, wo, pw, ho, ph, c) -> (bs, wo, ho, c, pw, ph)
+        inputs = tf.reshape(inputs, (-1, wo, ho, c_in * pw * ph))
+
+        if self.data_format == "channels_first":
+            inputs = tf.transpose(
+                inputs, [0, 3, 1, 2]  # (bs, w, h, c*pw*ph) -> (bs, c*pw*ph, w, h) ->
             )
+        return inputs
 
     def get_config(self):
         config = {
@@ -453,16 +466,29 @@ def __init__(
 
 
         Args:
-            pool_size: tuple describing the pool shape (p_h, p_w)
+            pool_size: integer or tuple describing the pool shape (p_h, p_w)
             data_format: can either be `channels_last` or `channels_first`
             name: name of the layer
             dtype: dtype of the layer
             **kwargs: params passed to the Layers constructor
         """
         super(InvertibleUpSampling, self).__init__(name=name, dtype=dtype, **kwargs)
-        self.pool_size = pool_size
         self.data_format = data_format
 
+        ndims = 2
+        ks: Tuple[int, ...]
+        if isinstance(pool_size, int):
+            ks = (pool_size,) * ndims
+        else:
+            ks = tuple(pool_size)
+
+        if len(ks) != ndims:
+            raise ValueError(
+                f"Expected {ndims}-dimensional pool_size, but "
+                f"got {len(ks)}-dimensional instead"
+            )
+        self.pool_size = ks
+
     def call(self, inputs):
         if self.data_format == "channels_first":
             # convert to channels_first
@@ -472,12 +498,12 @@ def call(self, inputs):
         w, h, c_in = input_shape[1], input_shape[2], input_shape[3]
         pw, ph = self.pool_size
         c = c_in // (pw * ph)
-        inputs = tf.reshape(inputs, (-1, w, h, pw, ph, c))
+        inputs = tf.reshape(inputs, (-1, w, h, c, pw, ph))
         inputs = tf.transpose(
             tf.reshape(
                 tf.transpose(
-                    inputs, [0, 5, 2, 4, 1, 3]
-                ),  # (bs, w, h, pw, ph, c) -> (bs, c, w, pw, h, ph)
+                    inputs, [0, 3, 2, 5, 1, 4]
+                ),  # (bs, w, h, c, pw, ph) -> (bs, c, w, pw, h, ph)
                 (-1, c, w, pw, h * ph),
             ),  # (bs, c, w, pw, h, ph) -> (bs, c, w, pw, h*ph) merge last axes
             [