Repository: deel-ai/deel-lip (language: python)

Issue:
Multiclass hkr
code to perform multiclass extension of HKR loss.
+ update of the binary classification to support (1,0) labels


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/python-linters.yml b/.github/workflows/python-linters.yml
index 54cd0e36..05cf5005 100644
--- a/.github/workflows/python-linters.yml
+++ b/.github/workflows/python-linters.yml
@@ -1,6 +1,6 @@
 name: deel-lip linters
 
-on: [push]
+on: [push, pull_request]
 
 jobs:
   checks:
@@ -8,7 +8,7 @@ jobs:
     strategy:
       max-parallel: 4
       matrix:
-        python-version: [3.6]
+        python-version: [3.6, 3.7, 3.8]
 
     steps:
     - uses: actions/checkout@v1
@@ -21,4 +21,4 @@ jobs:
         python -m pip install --upgrade pip
         pip install tox
     - name: Test with tox
-      run: tox -e py36-lint
\ No newline at end of file
+      run: tox -e py$(echo ${{ matrix.python-version }} | tr -d .)-lint
diff --git a/README.md b/README.md
index f1e8581b..271dda36 100644
--- a/README.md
+++ b/README.md
@@ -8,16 +8,16 @@
 Controlling the Lipschitz constant of a layer or a whole neural network has many applications ranging
 from adversarial robustness to Wasserstein distance estimation.
 
-This library provides implementation of **k-Lispchitz layers for `keras`**.
+This library provides an efficient implementation of **k-Lispchitz layers for `keras`**.
 
 ## The library contains:
 
  * k-Lipschitz variant of keras layers such as `Dense`, `Conv2D` and `Pooling`,
  * activation functions compatible with `keras`,
  * kernel initializers and kernel constraints for `keras`,
- * loss functions when working with Wasserstein distance estimations,
+ * loss functions that make use of Lipschitz constrained networks (see [our paper](https://arxiv.org/abs/2006.06520) for more information),
  * tools to monitor the singular values of kernels during training,
- * tools to convert k-Lipschitz network to regular network for faster evaluation.
+ * tools to convert k-Lipschitz network to regular network for faster inference.
 
 ## Example and usage
 
@@ -27,58 +27,98 @@ In order to make things simple the following rules have been followed during dev
 * When a k-Lipschitz layer overrides a standard keras layer, it uses the same interface and the same parameters.
   The only difference is a new parameter to control the Lipschitz constant of a layer.
 
-Here is a simple example showing how to build a 1-Lipschitz network:
+Here is an example showing how to build and train a 1-Lipschitz network:
 ```python
-from deel.lip.initializers import BjorckInitializer
-from deel.lip.layers import SpectralDense, SpectralConv2D
+from deel.lip.layers import SpectralDense, SpectralConv2D, ScaledL2NormPooling2D
 from deel.lip.model import Sequential
-from deel.lip.activations import PReLUlip
-from tensorflow.keras.layers import Input, Lambda, Flatten, MaxPool2D
-from tensorflow.keras import backend as K
+from deel.lip.activations import GroupSort
+from deel.lip.losses import HKR_multiclass_loss
+from tensorflow.keras.layers import Input, Flatten
 from tensorflow.keras.optimizers import Adam
+from tensorflow.keras.datasets import mnist
+from tensorflow.keras.utils import to_categorical
+import numpy as np
 
-# Sequential (resp Model) from deel.model has the same properties as any lipschitz
-# layer ( condense, setting of the lipschitz factor etc...). It act only as a container.
+# Sequential (resp Model) from deel.model has the same properties as any lipschitz model.
+# It act only as a container, with features specific to lipschitz 
+# functions (condensation, vanilla_exportation...)
 model = Sequential(
     [
-        Input(shape=(28, 28)),
-        Lambda(lambda x: K.reshape(x, (-1, 28, 28, 1))),
-
-        # Lipschitz layer preserve the API of their superclass ( here Conv2D )
+        Input(shape=(28, 28, 1)),
+        # Lipschitz layers preserve the API of their superclass ( here Conv2D )
         # an optional param is available: k_coef_lip which control the lipschitz
         # constant of the layer
         SpectralConv2D(
-            filters=32, kernel_size=(3, 3), padding='same',
-            activation=PReLUlip(), data_format='channels_last',
-            kernel_initializer=BjorckInitializer(15, 50)),
-        SpectralConv2D(
-            filters=32, kernel_size=(3, 3), padding='same',
-            activation=PReLUlip(), data_format='channels_last',
-            kernel_initializer=BjorckInitializer(15, 50)),
-        MaxPool2D(pool_size=(2, 2), data_format='channels_last'),
-
+            filters=16,
+            kernel_size=(3, 3),
+            activation=GroupSort(2),
+            use_bias=False,
+            kernel_initializer="orthogonal",
+        ),
+        # usual pooling layer are implemented (avg, max...), but new layers are also available
+        ScaledL2NormPooling2D(pool_size=(2, 2), data_format="channels_last"),
         SpectralConv2D(
-            filters=64, kernel_size=(3, 3), padding='same',
-            activation=PReLUlip(), data_format='channels_last',
-            kernel_initializer=BjorckInitializer(15, 50)),
-        SpectralConv2D(
-            filters=64, kernel_size=(3, 3), padding='same',
-            activation=PReLUlip(), data_format='channels_last',
-            kernel_initializer=BjorckInitializer(15, 50)),
-        MaxPool2D(pool_size=(2, 2), data_format='channels_last'),
-
+            filters=32,
+            kernel_size=(3, 3),
+            activation=GroupSort(2),
+            use_bias=False,
+            kernel_initializer="orthogonal",
+        ),
+        ScaledL2NormPooling2D(pool_size=(2, 2), data_format="channels_last"),
+        # our layers are fully interoperable with existing keras layers
         Flatten(),
-        SpectralDense(256, activation="relu", kernel_initializer=BjorckInitializer(15, 50)),
-        SpectralDense(10, activation="softmax"),
+        SpectralDense(
+            100,
+            activation=GroupSort(2),
+            use_bias=False,
+            kernel_initializer="orthogonal",
+        ),
+        SpectralDense(
+            10, activation=None, use_bias=False, kernel_initializer="orthogonal"
+        ),
     ],
-    k_coef_lip=0.5,
-    name='testing'
+    # similary model has a parameter to set the lipschitz constant
+    # to set automatically the constant of each layer
+    k_coef_lip=1.0,
+    name="hkr_model",
+)
+
+# HKR (Hinge-Krantorovich-Rubinstein) optimize robustness along with accuracy
+model.compile(
+    loss=HKR_multiclass_loss(alpha=5.0, min_margin=0.5),
+    optimizer=Adam(lr=0.01),
+    metrics=["accuracy"],
+)
+
+model.summary()
+
+# load data
+(x_train, y_train), (x_test, y_test) = mnist.load_data()
+# standardize and reshape the data
+x_train = np.expand_dims(x_train, -1)
+mean = x_train.mean()
+std = x_train.std()
+x_train = (x_train - mean) / std
+x_test = np.expand_dims(x_test, -1)
+x_test = (x_test - mean) / std
+# one hot encode the labels
+y_train = to_categorical(y_train)
+y_test = to_categorical(y_test)
+
+# fit the model
+model.fit(
+    x_train,
+    y_train,
+    batch_size=256,
+    epochs=15,
+    validation_data=(x_test, y_test),
+    shuffle=True,
 )
 
-optimizer = Adam(lr=0.001)
-model.compile(loss='categorical_crossentropy',
-              optimizer=optimizer,
-              metrics=['accuracy'])
+# once training is finished you can convert
+# SpectralDense layers into Dense layers and SpectralConv2D into Conv2D
+# which optimize performance for inference
+vanilla_model = model.vanilla_export()
 ```
 
 See [the full documentation](https://deel-lip.readthedocs.io) for a complete API description.
@@ -91,12 +131,13 @@ pip install deel-lip
 ```
 
 In order to use `deel-lip`, you also need a [valid tensorflow installation](https://www.tensorflow.org/install).
-`deel-lip` supports tensorflow from 2.0 to 2.2.
+`deel-lip` supports tensorflow versions 2.x
 
 ## Cite this work
 
 This library has been built to support the work presented in the paper
-*Achieving robustness in classification using optimaltransport with Hinge regularization*.
+[*Achieving robustness in classification using optimaltransport with Hinge regularization*](https://arxiv.org/abs/2006.06520)
+which aim provable and efficient robustness by design.
 
 This work can be cited as:
 ```latex
diff --git a/deel/lip/activations.py b/deel/lip/activations.py
index 835ad2af..3530facc 100644
--- a/deel/lip/activations.py
+++ b/deel/lip/activations.py
@@ -126,7 +126,6 @@ def build(self, input_shape):
             self.n = input_shape[self.channel_axis]
         if (input_shape[self.channel_axis] % self.n) != 0:
             raise RuntimeError("self.n has to be a divisor of the number of channels")
-        print(self.n)
 
     def _compute_lip_coef(self, input_shape=None):
         return 1.0
diff --git a/deel/lip/callbacks.py b/deel/lip/callbacks.py
index 05e137f1..df0f200f 100644
--- a/deel/lip/callbacks.py
+++ b/deel/lip/callbacks.py
@@ -34,7 +34,7 @@ def __init__(self, on_epoch: bool = True, on_batch: bool = False):
 
     def _condense_model(self):
         for layer in self.model.layers:
-            if isinstance(layer, Condensable):
+            if isinstance(layer, Condensable) or hasattr(layer, "condense"):
                 layer.condense()
 
     def on_train_batch_end(self, batch: int, logs: Optional[Dict[str, float]] = None):
diff --git a/deel/lip/layers.py b/deel/lip/layers.py
index d8ba7666..dafbfa3c 100644
--- a/deel/lip/layers.py
+++ b/deel/lip/layers.py
@@ -29,6 +29,7 @@
 from tensorflow.keras import backend as K
 from tensorflow.keras.initializers import RandomNormal
 from tensorflow.keras.layers import (
+    Layer,
     Dense,
     Conv2D,
     AveragePooling2D,
@@ -261,7 +262,6 @@ def build(self, input_shape):
     def _compute_lip_coef(self, input_shape=None):
         return 1.0  # this layer don't require a corrective factor
 
-    @tf.function
     def call(self, x, training=None):
         if training:
             W_bar, _u, sigma = spectral_normalization(
@@ -510,7 +510,6 @@ def _compute_lip_coef(self, input_shape=None):
             coefLip = np.sqrt((h * w) / ((k1 * ho - zl1 - zr1) * (k2 * wo - zl2 - zr2)))
         return coefLip
 
-    @tf.function
     def call(self, x, training=None):
         W_shape = self.kernel.shape
         if training:
@@ -632,7 +631,6 @@ def build(self, input_shape):
     def _compute_lip_coef(self, input_shape=None):
         return 1.0
 
-    @tf.function
     def call(self, x):
         W_bar = self.kernel / tf.norm(self.kernel) * self._get_coef()
         kernel = self.kernel
@@ -744,7 +742,6 @@ def build(self, input_shape):
     def _compute_lip_coef(self, input_shape=None):
         return SpectralConv2D._compute_lip_coef(self, input_shape)
 
-    @tf.function
     def call(self, x):
         W_bar = (self.kernel / tf.norm(self.kernel)) * self._get_coef()
         outputs = K.conv2d(
@@ -850,7 +847,6 @@ def build(self, input_shape):
     def _compute_lip_coef(self, input_shape=None):
         return np.sqrt(np.prod(np.asarray(self.pool_size)))
 
-    @tf.function
     def call(self, x, training=None):
         return super(AveragePooling2D, self).call(x) * self._get_coef()
 
@@ -935,7 +931,7 @@ def __init__(
         self._kwargs = kwargs
 
     def build(self, input_shape):
-        super(AveragePooling2D, self).build(input_shape)
+        super(ScaledL2NormPooling2D, self).build(input_shape)
         self._init_lip_coef(input_shape)
         self.built = True
 
@@ -943,18 +939,22 @@ def _compute_lip_coef(self, input_shape=None):
         return np.sqrt(np.prod(np.asarray(self.pool_size)))
 
     @staticmethod
-    @tf.custom_gradient
-    def _sqrt(x, eps_grad_sqrt):
-        def grad(dy):
-            return dy / (2 * tf.sqrt(x + eps_grad_sqrt))
+    def _sqrt(eps_grad_sqrt):
+        @tf.custom_gradient
+        def sqrt_op(x):
+            sqrtx = tf.sqrt(x)
 
-        return tf.sqrt(x), grad
+            def grad(dy):
+                return dy / (2 * (sqrtx + eps_grad_sqrt))
+
+            return sqrtx, grad
+
+        return sqrt_op
 
-    @tf.function
     def call(self, x, training=None):
         return (
-            ScaledL2NormPooling2D._sqrt(
-                super(AveragePooling2D, self).call(tf.square(x)), self.eps_grad_sqrt
+            ScaledL2NormPooling2D._sqrt(self.eps_grad_sqrt)(
+                super(ScaledL2NormPooling2D, self).call(tf.square(x))
             )
             * self._get_coef()
         )
@@ -963,7 +963,7 @@ def get_config(self):
         config = {
             "k_coef_lip": self.k_coef_lip,
         }
-        base_config = super(AveragePooling2D, self).get_config()
+        base_config = super(ScaledL2NormPooling2D, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
 
 
@@ -1016,7 +1016,6 @@ def _compute_lip_coef(self, input_shape=None):
             raise RuntimeError("data format not understood: %s" % self.data_format)
         return lip_coef
 
-    @tf.function
     def call(self, x, training=None):
         return super(ScaledGlobalAveragePooling2D, self).call(x) * self._get_coef()
 
@@ -1026,3 +1025,151 @@ def get_config(self):
         }
         base_config = super(ScaledGlobalAveragePooling2D, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
+
+
+@_deel_export
+class InvertibleDownSampling(Layer):
+    def __init__(
+        self, pool_size, data_format="channels_last", name=None, dtype=None, **kwargs
+    ):
+        """
+
+        This pooling layer perform a reshape on the spacial dimensions: it take a
+        (bs, h, w, c) ( if channels_last ) and reshape it to a
+        (bs, h/p_h, w/p_w, c*p_w*p_h ), where p_w and p_h are the shape of the pool.
+        By doing this the image size is reduced while the number of channels is
+        increased.
+
+        References:
+            Anil et al. https://arxiv.org/abs/1911.00937
+
+        Note:
+            The image shape must be divisible by the pool shape.
+
+        Args:
+            pool_size: tuple describing the pool shape
+            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 build(self, input_shape):
+        return super(InvertibleDownSampling, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+        # inputs = super(InvertibleDownSampling, self).call(inputs, **kwargs)
+        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,
+            )
+        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,
+            )
+
+    def get_config(self):
+        config = {
+            "data_format": self.data_format,
+            "pool_size": self.pool_size,
+        }
+        base_config = super(InvertibleDownSampling, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+@_deel_export
+class InvertibleUpSampling(Layer):
+    def __init__(
+        self, pool_size, data_format="channels_last", name=None, dtype=None, **kwargs
+    ):
+        """
+
+        This Layer is the inverse of the InvertibleDownSampling layer. It take a
+        (bs, h, w, c) ( if channels_last ) and reshape it to a
+        (bs, h/p_h, w/p_w, c*p_w*p_h ), where p_w and p_h are the shape of the
+        pool. By doing this the image size is reduced while the number of
+        channels is increased.
+
+        References:
+            Anil et al. https://arxiv.org/abs/1911.00937
+
+        Note:
+            The input number of channels must be divisible by the `p_w*p_h`.
+
+
+        Args:
+            pool_size: 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
+
+    def build(self, input_shape):
+        return super(InvertibleUpSampling, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+        if self.data_format == "channels_first":
+            # convert to channels_first
+            inputs = tf.transpose(inputs, [0, 2, 3, 1])
+        # from shape (bs, w, h, c*pw*ph) to (bs, w, h, pw, ph, c)
+        bs, w, h = inputs.shape[:-1]
+        (
+            pw,
+            ph,
+        ) = self.pool_size
+        c = inputs.shape[-1] // (pw * ph)
+        print(c)
+        inputs = tf.reshape(inputs, (-1, w, h, pw, ph, c))
+        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)
+                (-1, c, w, pw, h * ph),
+            ),  # (bs, c, w, pw, h, ph) -> (bs, c, w, pw, h*ph) merge last axes
+            [
+                0,
+                1,
+                4,
+                2,
+                3,
+            ],  # (bs, c, w, pw, h*ph) -> (bs, c, h*ph, w, pw)
+            # put each axis back in place
+        )
+        inputs = tf.reshape(
+            inputs, (-1, c, h * ph, w * pw)
+        )  # (bs, c, h*ph, w, pw) -> (bs, c, h*ph, w*pw)
+        if self.data_format == "channels_last":
+            inputs = tf.transpose(
+                inputs, [0, 2, 3, 1]  # (bs, c, h*ph, w*pw) -> (bs, w*pw, h*ph, c)
+            )
+        return inputs
+
+    def get_config(self):
+        config = {
+            "data_format": self.data_format,
+            "pool_size": self.pool_size,
+        }
+        base_config = super(InvertibleUpSampling, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
diff --git a/deel/lip/losses.py b/deel/lip/losses.py
index b7986d45..9c7f29dc 100644
--- a/deel/lip/losses.py
+++ b/deel/lip/losses.py
@@ -3,7 +3,8 @@
 # CRIAQ and ANITI - https://www.deel.ai/
 # =====================================================================================
 """
-This module contains losses used in wasserstein distance estimation.
+This module contains losses used in wasserstein distance estimation. See
+https://arxiv.org/abs/2006.06520 for more information.
 """
 import numpy as np
 import tensorflow as tf
@@ -12,12 +13,20 @@
 
 
 @_deel_export
-def KR_loss(true_values=(0, 1)):
-    """
+def KR_loss():
+    r"""
     Loss to estimate wasserstein-1 distance using Kantorovich-Rubinstein duality.
+    The Kantorovich-Rubinstein duality is formulated as following:
 
-    Args:
-        true_values: tuple containing the two label for each predicted class
+    .. math::
+        W_1(\mu, \nu) =
+        \sup_{f \in Lip_1(\Omega)} \underset{\textbf{x} \sim \mu}{\mathbb{E}}
+        \left[f(\textbf{x} )\right] -
+        \underset{\textbf{x}  \sim \nu}{\mathbb{E}} \left[f(\textbf{x} )\right]
+
+
+    Where mu and nu stands for the two distributions, the distribution where the
+    label is 1 and the rest.
 
     Returns:
         Callable, the function to compute Wasserstein loss
@@ -26,21 +35,20 @@ def KR_loss(true_values=(0, 1)):
 
     @tf.function
     def KR_loss_fct(y_true, y_pred):
-        S0 = tf.equal(y_true, true_values[0])
-        S1 = tf.equal(y_true, true_values[1])
+        # create two boolean masks each selecting one distribution
+        S0 = tf.equal(y_true, 1)
+        S1 = tf.not_equal(y_true, 1)
+        # compute the KR dual representation
         return K.mean(tf.boolean_mask(y_pred, S0)) - K.mean(tf.boolean_mask(y_pred, S1))
 
     return KR_loss_fct
 
 
 @_deel_export
-def neg_KR_loss(true_values=(1, -1)):
-    """
+def neg_KR_loss():
+    r"""
     Loss to compute the negative wasserstein-1 distance using Kantorovich-Rubinstein
-    duality.
-
-    Args:
-        true_values: tuple containing the two label for each predicted class
+    duality. This allows the maximisation of the term using conventional optimizer.
 
     Returns:
         Callable, the function to compute negative Wasserstein loss
@@ -49,20 +57,29 @@ def neg_KR_loss(true_values=(1, -1)):
 
     @tf.function
     def neg_KR_loss_fct(y_true, y_pred):
-        return -KR_loss(true_values)(y_true, y_pred)
+        return -KR_loss()(y_true, y_pred)
 
     return neg_KR_loss_fct
 
 
 @_deel_export
-def HKR_loss(alpha, min_margin=1, true_values=(1, -1)):
-    """
+def HKR_loss(alpha, min_margin=1):
+    r"""
     Wasserstein loss with a regularization param based on hinge loss.
 
+    .. math::
+        \inf_{f \in Lip_1(\Omega)} \underset{\textbf{x} \sim P_-}{\mathbb{E}}
+        \left[f(\textbf{x} )\right] - \underset{\textbf{x}  \sim P_+}
+        {\mathbb{E}} \left[f(\textbf{x} )\right] + \alpha
+        \underset{\textbf{x}}{\mathbb{E}} \left(\text{min_margin}
+        -Yf(\textbf{x})\right)_+
+
     Args:
         alpha: regularization factor
         min_margin: minimal margin ( see hinge_margin_loss )
-        true_values: tuple containing the two label for each predicted class
+        Kantorovich-rubinstein term of the loss. In order to be consistent
+        between hinge and KR, the first label must yield the positve class
+        while the second yields negative class.
 
     Returns:
          a function that compute the regularized Wasserstein loss
@@ -71,35 +88,169 @@ def HKR_loss(alpha, min_margin=1, true_values=(1, -1)):
 
     @tf.function
     def HKR_loss_fct(y_true, y_pred):
-        if alpha == np.inf:  # alpha negative hinge only
+        if alpha == np.inf:  # if alpha infinite, use hinge only
             return hinge_margin_loss(min_margin)(y_true, y_pred)
         else:
             # true value: positive value should be the first to be coherent with the
             # hinge loss (positive y_pred)
-            return alpha * hinge_margin_loss(min_margin)(y_true, y_pred) - KR_loss(
-                true_values
-            )(y_true, y_pred)
+            return alpha * hinge_margin_loss(min_margin)(y_true, y_pred) - KR_loss()(
+                y_true, y_pred
+            )
 
     return HKR_loss_fct
 
 
 @_deel_export
 def hinge_margin_loss(min_margin=1):
-    """
+    r"""
     Compute the hinge margin loss.
 
+    .. math::
+        \underset{\textbf{x}}{\mathbb{E}} \left(\text{min_margin}
+        -Yf(\textbf{x})\right)_+
+
     Args:
         min_margin: the minimal margin to enforce.
 
     Returns:
-        a function that compute the hinge loss
+        a function that compute the hinge loss.
 
     """
+    eps = 1e-7
 
     @tf.function
     def hinge_margin_fct(y_true, y_pred):
-        sign = K.sign(y_true)
+        sign = K.sign(
+            y_true - eps
+        )  # subtracting a small eps makes the loss work for (1,0) and (1,-1) labels
         hinge = K.maximum(0.0, min_margin - sign * y_pred)
         return K.mean(hinge)
 
     return hinge_margin_fct
+
+
+@_deel_export
+def KR_multiclass_loss():
+    r"""
+    Loss to estimate average of W1 distance using Kantorovich-Rubinstein
+    duality over outputs. In this multiclass setup thr KR term is computed
+    for each class and then averaged.
+
+    Note:
+        y_true has to be one hot encoded (labels being 1s and 0s ).
+
+    Returns:
+        Callable, the function to compute Wasserstein multiclass loss.
+
+    """
+
+    @tf.function
+    def KR_multiclass_loss_fct(y_true, y_pred):
+        # use y_true to zero out y_pred where y_true != 1
+        # espYtrue is the avg value of y_pred when y_true==1
+        # (one average per output neuron)
+        espYtrue = tf.reduce_sum(y_pred * y_true, axis=0) / tf.reduce_sum(
+            y_true, axis=0
+        )
+        # use(1- y_true) to zero out y_pred where y_true == 1
+        # espNotYtrue is the avg value of y_pred when y_true==0
+        # (one average per output neuron)
+        espNotYtrue = tf.reduce_sum(y_pred * (1 - y_true), axis=0) / (
+            tf.cast(tf.shape(y_true)[0], dtype="float32")
+            - tf.reduce_sum(y_true, axis=0)
+        )
+        # compute the differences to have the KR term for each output neuron,
+        # and compute the average over the classes
+        return tf.reduce_mean(-espNotYtrue + espYtrue)
+
+    return KR_multiclass_loss_fct
+
+
+@_deel_export
+def Hinge_multiclass_loss(min_margin=1):
+    """
+    Loss to estimate the Hinge loss in a multiclass setup. It compute the
+    elementwise hinge term. Note that this formulation differs from the
+    one commonly found in tensorflow/pytorch (with marximise the difference
+    between the two largest logits). This formulation is consistent with the
+    binary cassification loss used in a multiclass fashion.
+
+    Note:
+         y_true should be one hot encoded. labels in (1,0)
+
+    Returns:
+        Callable, the function to compute multiclass Hinge loss
+
+    """
+
+    @tf.function
+    def Hinge_multiclass_loss_fct(y_true, y_pred):
+        # convert (1,0) labels into (1,-1)
+        sign = 2 * y_true - 1
+        # compute the elementwise hinge term
+        hinge = tf.nn.relu(min_margin - sign * y_pred)
+        return K.mean(hinge)
+
+    return Hinge_multiclass_loss_fct
+
+
+@_deel_export
+def HKR_multiclass_loss(alpha=0.0, min_margin=1):
+    """
+    The multiclass version of HKR. This is done by computing the HKR term
+     over each class and averaging the results.
+
+    Args:
+        alpha: regularization factor
+        min_margin: minimal margin ( see Hinge_multiclass_loss )
+
+    Note:
+        y_true has to be one hot encoded.
+
+    Returns:
+        Callable, the function to compute HKR loss
+
+    """
+    hingeloss = Hinge_multiclass_loss(min_margin)
+    KRloss = KR_multiclass_loss()
+
+    @tf.function
+    def HKR_multiclass_loss_fct(y_true, y_pred):
+        if alpha == np.inf:  # alpha = inf => hinge only
+            return hingeloss(y_true, y_pred)
+        elif alpha == 0.0:  # alpha = 0 => KR only
+            return -KRloss(y_true, y_pred)
+        else:
+            return -KRloss(y_true, y_pred) / alpha + hingeloss(y_true, y_pred)
+
+    return HKR_multiclass_loss_fct
+
+
+@_deel_export
+def MultiMarginLoss(min_margin=1):
+    """
+    Compute the mean hinge margin loss for multi class (equivalent to Pytorch
+     multi_margin_loss)
+
+    Args:
+        min_margin: the minimal margin to enforce.
+        y_true has to be to_categorical
+
+    Returns:
+        Callable, the function to compute multi margin loss
+    """
+
+    @tf.function
+    def MultiMargin_fct(y_true, y_pred):
+        # get the y_pred[target_class]
+        # (zeroing out all elements of y_pred where y_true=0)
+        vYtrue = tf.reduce_sum(y_pred * y_true, axis=1, keepdims=True)
+        # computing elementwise margin term : margin + y_pred[i]-y_pred[target_class]
+        margin = tf.nn.relu(min_margin - vYtrue + y_pred)
+        # averaging on all outputs and batch
+        final_loss = tf.reduce_mean(
+            tf.where(y_true == 1, 0.0, margin)
+        )  # two steps is useless
+        return final_loss
+
+    return MultiMargin_fct
diff --git a/deel/lip/model.py b/deel/lip/model.py
index feb693aa..a537f361 100644
--- a/deel/lip/model.py
+++ b/deel/lip/model.py
@@ -7,10 +7,12 @@
 for condensation and vanilla exportation.
 """
 import math
+from typing import Dict
 from warnings import warn
 import numpy as np
+from tensorflow import Tensor
 from tensorflow.keras import Sequential as KerasSequential, Model as KerasModel
-from tensorflow.keras.layers import Input
+from tensorflow.keras.layers import Input, InputLayer
 from .layers import LipschitzLayer, Condensable
 from .utils import _deel_export
 
@@ -35,7 +37,7 @@ def __init__(
             name: name of the model, can be None
             k_coef_lip: the Lipschitz coefficient to ensure globally on the model.
         """
-        super().__init__(layers, name)
+        super(Sequential, self).__init__(layers, name)
         self.set_klip_factor(k_coef_lip)
 
     def build(self, input_shape=None):
@@ -117,47 +119,82 @@ def get_config(self):
 
 
 @_deel_export
-class Model(KerasModel, Condensable):
-    def __init__(self, *args, **kwargs):
-        """
-        Equivalent of keras.Model but support condensation and vanilla exportation.
+class Model(KerasModel):
+    """
+    Equivalent of keras.Model but support condensation and vanilla exportation.
 
-        Warning:
-             As lipschitz constant are multiplicative along layer, the Model class
-             cannot set a global Lipschitz constant (problem with branching inside a
-             model).
-        """
-        super().__init__(*args, **kwargs)
+    Warning:
+         As lipschitz constant are multiplicative along layer, the Model class
+         cannot set a global Lipschitz constant (problem with branching inside a
+         model).
+    """
 
     def condense(self):
         for layer in self.layers:
             if isinstance(layer, Condensable):
                 layer.condense()
 
-    def vanilla_export(self):
-        dict_src = {}
-        input_l = Input(self.input.shape[1:])
-        mdl = input_l
-        dict_src[self.input.name] = mdl
+    def vanilla_export(self) -> KerasModel:
+        """
+        Export this model to a "Vanilla" model, i.e. a model without Condensable
+        layers.
+
+        Returns:
+            A Keras model, identical to this model, but where condensable layers have
+            been replaced with their vanilla equivalent (e.g. SpectralConv2D with
+             Conv2D).
+        """
+        # Dictionary that will map tensor names (from the current model) to tensors
+        # in the exported model.# We initialize the dictionary for inputs:
+        tensors: Dict[str, Tensor] = {}
+
+        # Initialize the dictionary with inputs:
+        tensors.update({inp.name: Input(shape=inp.shape[1:]) for inp in self.inputs})
+
         for lay in self.layers:
-            if lay.input.name in dict_src:
-                print("tensor input shape " + str(dict_src[lay.input.name].shape))
-            if lay.input.name == lay.output.name:
-                pass
+
+            # Skip input layers:
+            if isinstance(lay, InputLayer):
+                continue
+
+            # Condense+Export the layer if it is a non-vanilla layer, otherwise
+            # just copy the layer:
+            if isinstance(lay, Condensable):
+                lay.condense()
+                lay_cp = lay.vanilla_export()
             else:
-                if isinstance(lay, Condensable):
-                    lay.condense()
-                    new_lay = lay.vanilla_export()
-                    mdl = new_lay(dict_src[lay.input.name])
+                # Duplicate layer (weights are not duplicated):
+                lay_cp = lay.__class__.from_config(lay.get_config())
+                lay_cp.build(lay.input_shape)
+                lay_cp.set_weights(lay.get_weights().copy())
+
+            # For each input nodes, we are going to create corresponding operations
+            # in the exported models:
+            for inode in range(len(lay.inbound_nodes)):
+                inputs = lay.get_input_at(inode)
+                outputs = lay.get_output_at(inode)
+
+                # Fetch the
+                if isinstance(inputs, list):
+                    inputs = [tensors[input.name] for input in inputs]
                 else:
-                    # duplicate layer (warning weights are not duplicated)
-                    lay_cp = lay.__class__.from_config(lay.get_config())
-                    lay_cp.build(lay.input.shape)
-                    lay_cp.set_weights(lay.get_weights().copy())
-                    mdl = lay_cp(dict_src[lay.input.name])
-            dict_src[lay.output.name] = mdl
-        net = KerasModel(input_l, mdl)
-        return net
+                    inputs = tensors[inputs.name]
+
+                # Retrieve outputs layers (for the exported layer):
+                moutputs = lay_cp(inputs)
+
+                # Add the output tensors to the dictionary, using the names from the
+                # original model:
+                if isinstance(outputs, list):
+                    for outi, mouti in zip(outputs, moutputs):
+                        tensors[outi.name] = mouti
+                else:
+                    tensors[outputs.name] = moutputs
+
+        return KerasModel(
+            [tensors[inp.name] for inp in self.inputs],
+            [tensors[out.name] for out in self.outputs],
+        )
 
 
 vanillaModel = Model.vanilla_export
diff --git a/deel/lip/normalizers.py b/deel/lip/normalizers.py
index f052fa56..3a1e005b 100644
--- a/deel/lip/normalizers.py
+++ b/deel/lip/normalizers.py
@@ -6,6 +6,7 @@
 This module contains computation function, for BjorckNormalizer and spectral
 normalization. This is done for internal use only.
 """
+import tensorflow as tf
 from tensorflow.keras import backend as K
 
 DEFAULT_NITER_BJORCK = 15
@@ -13,6 +14,7 @@
 DEFAULT_NITER_SPECTRAL_INIT = 10
 
 
+@tf.function
 def bjorck_normalization(w, niter=DEFAULT_NITER_BJORCK):
     """
     apply Bjorck normalization on w.
@@ -52,6 +54,7 @@ def _power_iteration(w, u, niter=DEFAULT_NITER_SPECTRAL):
     return _u, _v
 
 
+@tf.function
 def spectral_normalization(kernel, u=None, niter=DEFAULT_NITER_SPECTRAL):
     """
     Normalize the kernel to have it's max eigenvalue == 1.
diff --git a/doc/source/deel.lip.layers.rst b/doc/source/deel.lip.layers.rst
deleted file mode 100644
index 48b7b858..00000000
--- a/doc/source/deel.lip.layers.rst
+++ /dev/null
@@ -1,9 +0,0 @@
-.. _deel-lip-api-layers:
-
-deel.lip.layers module
-======================
-
-.. automodule:: deel.lip.layers
-   :members:
-   :undoc-members:
-   :show-inheritance:
diff --git a/doc/source/deel.lip.rst b/doc/source/deel.lip.rst
deleted file mode 100644
index ae1054ef..00000000
--- a/doc/source/deel.lip.rst
+++ /dev/null
@@ -1,28 +0,0 @@
-.. _deel-lip-api:
-
-deel.lip package
-================
-
-Submodules
-----------
-
-.. toctree::
-   :maxdepth: 4
-
-   deel.lip.activations
-   deel.lip.callbacks
-   deel.lip.constraints
-   deel.lip.initializers
-   deel.lip.layers
-   deel.lip.losses
-   deel.lip.model
-   deel.lip.normalizers
-   deel.lip.utils
-
-Module contents
----------------
-
-.. automodule:: deel.lip
-   :members:
-   :undoc-members:
-   :show-inheritance:
diff --git a/doc/source/index.rst b/doc/source/index.rst
index d07b084b..8ac93f6f 100644
--- a/doc/source/index.rst
+++ b/doc/source/index.rst
@@ -5,61 +5,188 @@
 
 Welcome to deel-lip documentation!
 ==================================================
+Controlling the Lipschitz constant of a layer or a whole neural network
+has many applications ranging from adversarial robustness to Wasserstein
+distance estimation.
 
-Controlling the Lipschitz constant of a layer or a whole neural network has many applications ranging
-from adversarial robustness to Wasserstein distance estimation.
-
-This library provides implementation of **k-Lispchitz layers for** ``keras``.
+This library provides an efficient implementation of **k-Lispchitz
+layers for keras**.
 
 The library contains:
 ---------------------
 
- * k-Lipschitz variant of keras layers such as `Dense`, `Conv2D` and `Pooling`,
- * activation functions compatible with `keras`,
- * kernel initializers and kernel constraints for `keras`,
- * loss functions when working with Wasserstein distance estimations,
- * tools to monitor the singular values of kernels during training,
- * tools to convert k-Lipschitz network to regular network for faster evaluation.
-
+-  k-Lipschitz variant of keras layers such as ``Dense``, ``Conv2D`` and
+   ``Pooling``,
+-  activation functions compatible with ``keras``,
+-  kernel initializers and kernel constraints for ``keras``,
+-  loss functions that make use of Lipschitz constrained networks (see
+   `our paper <https://arxiv.org/abs/2006.06520>`__ for more
+   information),
+-  tools to monitor the singular values of kernels during training,
+-  tools to convert k-Lipschitz network to regular network for faster
+   inference.
+
+Example and usage
+-----------------
+
+In order to make things simple the following rules have been followed
+during development:
+- ``deel-lip`` follows the ``keras`` package structure.
+- All elements (layers, activations, initializers, ...) are compatible with standard the ``keras`` elements.
+- When a k-Lipschitz layer overrides a standard keras layer, it uses the same interface and
+the same parameters. The only difference is a new parameter to control the Lipschitz constant of a layer.
+
+Here is an example showing how to build and train a 1-Lipschitz network:
+
+.. code:: python
+
+    from deel.lip.layers import SpectralDense, SpectralConv2D, ScaledL2NormPooling2D
+    from deel.lip.model import Sequential
+    from deel.lip.activations import GroupSort
+    from deel.lip.losses import HKR_multiclass_loss
+    from tensorflow.keras.layers import Input, Flatten
+    from tensorflow.keras.optimizers import Adam
+    from tensorflow.keras.datasets import mnist
+    from tensorflow.keras.utils import to_categorical
+    import numpy as np
+
+    # Sequential (resp Model) from deel.model has the same properties as any lipschitz model.
+    # It act only as a container, with features specific to lipschitz
+    # functions (condensation, vanilla_exportation...)
+    model = Sequential(
+        [
+            Input(shape=(28, 28, 1)),
+            # Lipschitz layers preserve the API of their superclass ( here Conv2D )
+            # an optional param is available: k_coef_lip which control the lipschitz
+            # constant of the layer
+            SpectralConv2D(
+                filters=16,
+                kernel_size=(3, 3),
+                activation=GroupSort(2),
+                use_bias=False,
+                kernel_initializer="orthogonal",
+            ),
+            # usual pooling layer are implemented (avg, max...), but new layers are also available
+            ScaledL2NormPooling2D(pool_size=(2, 2), data_format="channels_last"),
+            SpectralConv2D(
+                filters=32,
+                kernel_size=(3, 3),
+                activation=GroupSort(2),
+                use_bias=False,
+                kernel_initializer="orthogonal",
+            ),
+            ScaledL2NormPooling2D(pool_size=(2, 2), data_format="channels_last"),
+            # our layers are fully interoperable with existing keras layers
+            Flatten(),
+            SpectralDense(
+                100,
+                activation=GroupSort(2),
+                use_bias=False,
+                kernel_initializer="orthogonal",
+            ),
+            SpectralDense(
+                10, activation=None, use_bias=False, kernel_initializer="orthogonal"
+            ),
+        ],
+        # similary model has a parameter to set the lipschitz constant
+        # to set automatically the constant of each layer
+        k_coef_lip=1.0,
+        name="hkr_model",
+    )
+
+    # HKR (Hinge-Krantorovich-Rubinstein) optimize robustness along with accuracy
+    model.compile(
+        loss=HKR_multiclass_loss(alpha=5.0, min_margin=0.5),
+        optimizer=Adam(lr=0.01),
+        metrics=["accuracy"],
+    )
+
+    model.summary()
+
+    # load data
+    (x_train, y_train), (x_test, y_test) = mnist.load_data()
+    # standardize and reshape the data
+    x_train = np.expand_dims(x_train, -1)
+    mean = x_train.mean()
+    std = x_train.std()
+    x_train = (x_train - mean) / std
+    x_test = np.expand_dims(x_test, -1)
+    x_test = (x_test - mean) / std
+    # one hot encode the labels
+    y_train = to_categorical(y_train)
+    y_test = to_categorical(y_test)
+
+    # fit the model
+    model.fit(
+        x_train,
+        y_train,
+        batch_size=256,
+        epochs=15,
+        validation_data=(x_test, y_test),
+        shuffle=True,
+    )
+
+    # once training is finished you can convert
+    # SpectralDense layers into Dense layers and SpectralConv2D into Conv2D
+    # which optimize performance for inference
+    vanilla_model = model.vanilla_export()
+
+See `the full documentation <https://deel-lip.readthedocs.io>`__ for a
+complete API description.
 
 Installation
 ------------
 
 You can install ``deel-lip`` directly from pypi:
 
-.. code-block:: bash
+.. code:: bash
 
-   pip install deel-lip
+    pip install deel-lip
 
-In order to use ``deel-lip``, you also need a `valid tensorflow installation <https://www.tensorflow.org/install>`_.
-``deel-lip`` supports tensorflow from 2.0 to 2.2.
+In order to use ``deel-lip``, you also need a `valid tensorflow
+installation <https://www.tensorflow.org/install>`__. ``deel-lip``
+supports tensorflow versions 2.x
 
 Cite this work
 --------------
 
-.. raw:: html
-
-   This library has been built to support the work presented in the paper
-   <a href="https://arxiv.org/abs/2006.06520"><i>Achieving robustness in classification using optimal transport with Hinge regularization</i></a>.
-   This work can be cited as:
-
-.. code-block:: latex
-
-   @misc{2006.06520,
-   Author = {
-      Mathieu Serrurier
-      and Franck Mamalet
-      and Alberto González-Sanz
-      and Thibaut Boissin
-      and Jean-Michel Loubes
-      and Eustasio del Barrio
-   },
-   Title = {
-      Achieving robustness in classification using optimal transport with hinge regularization
-   },
-   Year = {2020},
-   Eprint = {arXiv:2006.06520},
-   }
+This library has been built to support the work presented in the paper
+`Achieving robustness in classification using optimaltransport with
+Hinge regularization <https://arxiv.org/abs/2006.06520>`__ which aim
+provable and efficient robustness by design.
+
+This work can be cited as:
+
+.. code:: latex
+
+    @misc{2006.06520,
+    Author = {Mathieu Serrurier and Franck Mamalet and Alberto González-Sanz and Thibaut Boissin and Jean-Michel Loubes and Eustasio del Barrio},
+    Title = {Achieving robustness in classification using optimal transport with hinge regularization},
+    Year = {2020},
+    Eprint = {arXiv:2006.06520},
+    }
+
+Contributing
+------------
+
+To contribute, you can open an
+`issue <https://github.com/deel-ai/deel-lip/issues>`__, or fork this
+repository and then submit changes through a
+`pull-request <https://github.com/deel-ai/deel-lip/pulls>`__.
+We use `black <https://pypi.org/project/black/>`__ to format the code and follow PEP-8 convention.
+To check that your code will pass the lint-checks, you can run:
+
+.. code:: bash
+
+    tox -e py36-lint
+
+You need `tox <https://tox.readthedocs.io/en/latest/>`__ in order to
+run this. You can install it via ``pip``:
+
+.. code:: bash
+
+    pip install tox
+
 
 Indices and tables
 ==================
diff --git a/setup.cfg b/setup.cfg
index 793f9a75..e75e1925 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -7,14 +7,14 @@ extend-ignore =
 
 
 [tox:tox]
-envlist = py36,py37,py38,py36-lint
+envlist = py{36,37,38},py{36,37,38}-lint
 
 [testenv]
 deps =
 commands =
     python -m unittest
 
-[testenv:py36-lint]
+[testenv:py{36,37,38}-lint]
 skip_install = true
 deps =
     black
diff --git a/setup.py b/setup.py
index 383a0f88..cdb6848d 100644
--- a/setup.py
+++ b/setup.py
@@ -24,12 +24,13 @@
     "recommonmark",
     "sphinx_rtd_theme",
     "sphinx_markdown_builder",
+    "sphinxcontrib_katex",
     "ipython",  # required for Pygments
 ]
 
 setuptools.setup(
     name="deel-lip",
-    version="1.0.2",
+    version="1.1.0",
     author=", ".join(["Mathieu SERRURIER", "Franck MAMALET", "Thibaut BOISSIN"]),
     author_email=", ".join(
         [
@@ -43,7 +44,7 @@
     long_description_content_type="text/markdown",
     url="https://github.com/deel-ai/deel-lip",
     packages=setuptools.find_namespace_packages(include=["deel.*"]),
-    install_requires=["numpy", "tensorflow>=2,<=2.3"],
+    install_requires=["numpy", "tensorflow>=2,<3"],
     license="MIT",
     extras_require={"dev": dev_requires, "docs": docs_requires},
     classifiers=[


Test patch:
diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml
new file mode 100644
index 00000000..3f6bbcf4
--- /dev/null
+++ b/.github/workflows/python-tests.yml
@@ -0,0 +1,24 @@
+name: deel-lip tests
+
+on: [push, pull_request]
+
+jobs:
+  checks:
+    runs-on: ubuntu-latest
+    strategy:
+      max-parallel: 4
+      matrix:
+        python-version: [3.6, 3.7, 3.8]
+
+    steps:
+    - uses: actions/checkout@v1
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v1
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install tox
+    - name: Test with tox
+      run: tox -e py$(echo ${{ matrix.python-version }} | tr -d .)
\ No newline at end of file
diff --git a/tests/test_condense.py b/tests/test_condense.py
index 8b899a92..a3097062 100644
--- a/tests/test_condense.py
+++ b/tests/test_condense.py
@@ -13,6 +13,7 @@
     SpectralConv2D,
     FrobeniusDense,
     FrobeniusConv2D,
+    ScaledL2NormPooling2D,
 )
 import tensorflow as tf
 from tensorflow.keras.layers import Flatten, Dense
@@ -28,7 +29,7 @@
 
 class MyTestCase(unittest.TestCase):
     def test_as_supertype_sequential(self):
-        input_shape = (5, 5, 3)
+        input_shape = (8, 8, 3)
         model = Sequential(
             [
                 # the bug occurs only when using the "as_supertype_export" function
@@ -39,6 +40,7 @@ def test_as_supertype_sequential(self):
                 # - tensorflow 2.1 ( not 2.0 )
                 Input(input_shape),
                 SpectralConv2D(2, (3, 3)),
+                ScaledL2NormPooling2D((2, 2)),
                 FrobeniusConv2D(2, (3, 3)),
                 Flatten(),
                 Dense(4),
@@ -50,15 +52,16 @@ def test_as_supertype_sequential(self):
         self._test_model(model, input_shape)
 
     def test_as_supertype_model(self):
-        input_shape = (5, 5, 3)
+        input_shape = (8, 8, 3)
         inp = Input(input_shape)
         x = SpectralConv2D(2, (3, 3), k_coef_lip=2.0)(inp)
+        x = ScaledL2NormPooling2D((2, 2), k_coef_lip=2.0)(x)
         x = FrobeniusConv2D(2, (3, 3), k_coef_lip=2.0)(x)
         x = Flatten()(x)
         x = Dense(4)(x)
         x = SpectralDense(4, k_coef_lip=2.0)(x)
-        x = FrobeniusDense(2, k_coef_lip=2.0)(x)
-        model = Model(inputs=inp, outputs=x)
+        out = FrobeniusDense(2, k_coef_lip=2.0)(x)
+        model = Model(inputs=inp, outputs=out)
         self._test_model(model, input_shape)
 
     def _test_model(self, model, input_shape):
diff --git a/tests/test_layers.py b/tests/test_layers.py
index 807d389e..323dbca5 100644
--- a/tests/test_layers.py
+++ b/tests/test_layers.py
@@ -27,6 +27,8 @@
     ScaledAveragePooling2D,
     ScaledGlobalAveragePooling2D,
     ScaledL2NormPooling2D,
+    InvertibleDownSampling,
+    InvertibleUpSampling,
 )
 from deel.lip.model import Sequential
 from deel.lip.utils import load_model, evaluate_lip_const
@@ -135,7 +137,6 @@ def generate_k_lip_model(layer_type: type, layer_params: dict, input_shape, k):
     if issubclass(layer_type, LipschitzLayer):
         layer_params["k_coef_lip"] = k
     layer = layer_type(**layer_params)
-    assert isinstance(layer, LipschitzLayer) or isinstance(layer, Dense)
     assert isinstance(layer, Layer)
     # print(input_shape)
     # print(layer.compute_output_shape((32, ) + input_shape))
@@ -593,7 +594,7 @@ def test_scaledl2normpooling2d(self):
                     layer_params={
                         "layers": [
                             Input((5, 5, 1)),
-                            ScaledL2NormPooling2D(data_format="channels_last"),
+                            ScaledL2NormPooling2D((2, 3), data_format="channels_last"),
                         ]
                     },
                     batch_size=1000,
@@ -609,7 +610,7 @@ def test_scaledl2normpooling2d(self):
                     layer_params={
                         "layers": [
                             Input((5, 5, 1)),
-                            ScaledL2NormPooling2D(data_format="channels_last"),
+                            ScaledL2NormPooling2D((2, 3), data_format="channels_last"),
                         ]
                     },
                     batch_size=1000,
@@ -625,7 +626,7 @@ def test_scaledl2normpooling2d(self):
                     layer_params={
                         "layers": [
                             Input((5, 5, 1)),
-                            ScaledL2NormPooling2D(data_format="channels_last"),
+                            ScaledL2NormPooling2D((2, 3), data_format="channels_last"),
                         ]
                     },
                     batch_size=1000,
@@ -747,6 +748,42 @@ def test_callbacks(self):
             ]
         )
 
+    def test_invertibledownsampling(self):
+        # tests only checks that lip cons is enforced
+        self._apply_tests_bank(
+            [
+                dict(
+                    layer_type=InvertibleDownSampling,
+                    layer_params={"pool_size": (2, 3)},
+                    batch_size=1000,
+                    steps_per_epoch=1,
+                    epochs=5,
+                    input_shape=(6, 6, 3),
+                    k_lip_data=1.0,
+                    k_lip_model=1.0,
+                    callbacks=[],
+                ),
+            ]
+        )
+
+    def test_invertibleupsampling(self):
+        # tests only checks that lip cons is enforced
+        self._apply_tests_bank(
+            [
+                dict(
+                    layer_type=InvertibleUpSampling,
+                    layer_params={"pool_size": (2, 3)},
+                    batch_size=1000,
+                    steps_per_epoch=1,
+                    epochs=5,
+                    input_shape=(6, 6, 18),
+                    k_lip_data=1.0,
+                    k_lip_model=1.0,
+                    callbacks=[],
+                ),
+            ]
+        )
+
 
 if __name__ == "__main__":
     unittest.main()