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

Issue:
Bump version to 1.4.0



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/python-linters.yml b/.github/workflows/python-linters.yml
index 05cf5005..d18ad3b4 100644
--- a/.github/workflows/python-linters.yml
+++ b/.github/workflows/python-linters.yml
@@ -1,6 +1,11 @@
 name: deel-lip linters
 
-on: [push, pull_request]
+on:
+  push:
+    branches:
+      - master
+      - develop
+  pull_request:
 
 jobs:
   checks:
@@ -8,12 +13,12 @@ jobs:
     strategy:
       max-parallel: 4
       matrix:
-        python-version: [3.6, 3.7, 3.8]
+        python-version: [3.7, "3.10"]
 
     steps:
-    - uses: actions/checkout@v1
+    - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
index 03760727..82ad4abd 100644
--- a/.github/workflows/python-publish.yml
+++ b/.github/workflows/python-publish.yml
@@ -13,9 +13,9 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: '3.x'
     - name: Install dependencies
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index b52a1954..aa5b7909 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -27,18 +27,12 @@ Before opening a pull request, please make sure you check your code and you run
 unit tests:
 
 ```bash
-# Using make:
 $ make test
-
-# Or using directly tox in your development environment:
-$ tox
 ```
 
 This command will:
 - check your code with black PEP-8 formatter and flake8 linter.
-- run `unittest` on the `tests/` folder with Python 3.6, 3.7 and 3.8.
-> Note: If you do not have those 3 interpreters, the tests will only be only performed
-  with your current interpreter.
+- run `unittest` on the `tests/` folder with different Python and TensorFlow versions.
 
 
 ## Submitting your changes
diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 00000000..86a41807
--- /dev/null
+++ b/MANIFEST.in
@@ -0,0 +1,2 @@
+include deel/lip/VERSION
+include LICENSE
diff --git a/Makefile b/Makefile
index 95fa8365..d1adc7de 100644
--- a/Makefile
+++ b/Makefile
@@ -22,10 +22,16 @@ prepare-dev:
 	. deel_lip_dev_env/bin/activate && pip install -e .[docs]
 
 test:
-	. deel_lip_dev_env/bin/activate && tox
+	. deel_lip_dev_env/bin/activate && tox -e py37-tf23
+	. deel_lip_dev_env/bin/activate && tox -e py39-tf27
+	. deel_lip_dev_env/bin/activate && tox -e py310-tflatest
+	. deel_lip_dev_env/bin/activate && tox -e py310-lint
 
 test-disable-gpu:
-	. deel_lip_dev_env/bin/activate && CUDA_VISIBLE_DEVICES=-1 tox
+	. deel_lip_dev_env/bin/activate && CUDA_VISIBLE_DEVICES=-1 tox -e py37-tf23
+	. deel_lip_dev_env/bin/activate && CUDA_VISIBLE_DEVICES=-1 tox -e py39-tf27
+	. deel_lip_dev_env/bin/activate && CUDA_VISIBLE_DEVICES=-1 tox -e py310-tflatest
+	. deel_lip_dev_env/bin/activate && CUDA_VISIBLE_DEVICES=-1 tox -e py310-lint
 
 doc:
 	. deel_lip_dev_env/bin/activate && cd doc && make html && cd -
diff --git a/deel/lip/VERSION b/deel/lip/VERSION
new file mode 100644
index 00000000..88c5fb89
--- /dev/null
+++ b/deel/lip/VERSION
@@ -0,0 +1,1 @@
+1.4.0
diff --git a/deel/lip/__init__.py b/deel/lip/__init__.py
index aea6c3bd..007ae98a 100644
--- a/deel/lip/__init__.py
+++ b/deel/lip/__init__.py
@@ -2,6 +2,12 @@
 # rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
 # CRIAQ and ANITI - https://www.deel.ai/
 # =====================================================================================
+
+from os import path
+
+with open(path.join(path.dirname(__file__), "VERSION")) as f:
+    __version__ = f.read().strip()
+
 from . import activations
 from . import callbacks
 from . import constraints
@@ -12,4 +18,5 @@
 from . import metrics
 from .model import Sequential, Model, vanillaModel
 from . import normalizers
+from . import regularizers
 from . import utils
diff --git a/deel/lip/activations.py b/deel/lip/activations.py
index e0efe12e..8a459c7e 100644
--- a/deel/lip/activations.py
+++ b/deel/lip/activations.py
@@ -7,6 +7,7 @@
 It can be added as a layer, or it can be used in the "activation" params for other
 layers.
 """
+import math
 import tensorflow as tf
 from tensorflow.keras import backend as K
 from tensorflow.keras.constraints import MinMaxNorm
@@ -211,3 +212,85 @@ def PReLUlip(k_coef_lip=1.0):
     return PReLU(
         alpha_constraint=MinMaxNorm(min_value=-k_coef_lip, max_value=k_coef_lip)
     )
+
+
+@register_keras_serializable("deel-lip", "Householder")
+class Householder(Layer, LipschitzLayer):
+    def __init__(
+        self,
+        data_format="channels_last",
+        k_coef_lip=1.0,
+        theta_initializer=None,
+        **kwargs,
+    ):
+        """
+        Householder activation: https://openreview.net/pdf?id=tD7eCtaSkR
+        From https://github.com/singlasahil14/SOC
+
+        Args:
+            data_format: either channels_first or channels_last. Only channels_last is
+                supported.
+            k_coef_lip: The lipschitz coefficient to be enforced.
+            theta_initializer: initializer for the angle theta of reflection. Defaults
+                to pi/2, which corresponds to GroupSort2.
+            **kwargs: parameters passed to the `tf.keras.layers.Layer`.
+
+        Input shape:
+            Arbitrary. Use the keyword argument `input_shape` (tuple of integers, does
+            not include the samples axis) when using this layer as the first layer in a
+            model.
+
+        Output shape:
+            Same size as input.
+
+        """
+        if data_format != "channels_last":
+            raise RuntimeError("Only 'channels_last' data format is supported")
+
+        self.data_format = data_format
+        self.set_klip_factor(k_coef_lip)
+        self.theta_initializer = theta_initializer
+        super().__init__(**kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+        self._init_lip_coef(input_shape)
+        if (input_shape[-1] % 2) != 0:
+            raise RuntimeError("2 has to be a divisor of the number of channels")
+
+        self.theta = self.add_weight(
+            "theta",
+            shape=[input_shape[-1] // 2],
+            initializer=self.theta_initializer,
+        )
+        if self.theta_initializer is None:
+            self.theta.assign(tf.ones_like(self.theta, dtype=tf.float32) * math.pi / 2)
+
+    def _compute_lip_coef(self, input_shape=None):
+        return 1.0
+
+    def call(self, x):
+        z1, z2 = tf.split(x, 2, axis=-1)
+
+        # selector > 0 if point (z1, z2) is on one side of reflection line, else < 0.
+        # Reflection line is defined by angle theta/2.
+        selector = (z1 * tf.sin(0.5 * self.theta)) - (z2 * tf.cos(0.5 * self.theta))
+
+        cos_theta = tf.cos(self.theta)
+        sin_theta = tf.sin(self.theta)
+        reflected_z1 = z1 * cos_theta + z2 * sin_theta
+        reflected_z2 = z1 * sin_theta - z2 * cos_theta
+
+        a = tf.where(selector <= 0, z1, reflected_z1)
+        b = tf.where(selector <= 0, z2, reflected_z2)
+
+        return tf.concat([a, b], axis=-1)
+
+    def get_config(self):
+        config = {
+            "k_coef_lip": self.k_coef_lip,
+            "data_format": self.data_format,
+            "theta_initializer": self.theta_initializer,
+        }
+        base_config = super().get_config()
+        return dict(list(base_config.items()) + list(config.items()))
diff --git a/deel/lip/callbacks.py b/deel/lip/callbacks.py
index 459b32a1..b97e69c9 100644
--- a/deel/lip/callbacks.py
+++ b/deel/lip/callbacks.py
@@ -10,7 +10,7 @@
 
 import tensorflow as tf
 from tensorflow.keras.callbacks import Callback
-
+import numpy as np
 from .layers import Condensable
 
 
@@ -164,3 +164,65 @@ def get_config(self):
         }
         base_config = super(MonitorCallback, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
+
+
+class LossParamScheduler(Callback):
+    def __init__(self, param_name, fp, xp, step=0):
+        """
+        Scheduler to modify a loss parameter during training. It uses a linear
+        interpolation (defined by fp and xp) depending on the optimization step.
+
+        Args:
+            param_name (str): name of the parameter of the loss to tune. Must be a
+                tf.Variable.
+            fp (list): values of the loss parameter as steps given by the xp.
+            xp (list): step where the parameter equals fp.
+            step: step value, for serialization/deserialization purposes.
+        """
+        self.xp = xp
+        self.fp = fp
+        self.step = step
+        self.param_name = param_name
+
+    def on_train_batch_begin(self, batch: int, logs=None):
+        new_value = np.interp(self.step, self.xp, self.fp)
+        self.model.loss.__getattribute__(self.param_name).assign(new_value)
+        self.step += 1
+        super(LossParamScheduler, self).on_train_batch_end(batch, logs)
+
+    def get_config(self):
+        return {
+            "xp": self.xp,
+            "fp": self.fp,
+            "step": self.step,
+            "param_name": self.param_name,
+        }
+
+
+class LossParamLog(Callback):
+    def __init__(self, param_name, rate=1):
+        """
+        Logger to print values of a loss parameter at each epoch.
+
+        Args:
+            param_name (str): name of the parameter of the loss to log.
+            rate (int): logging rate (in epochs)
+        """
+        self.param_name = param_name
+        self.rate = rate
+
+    def on_epoch_end(self, epoch: int, logs=None):
+        if epoch % self.rate == 0:
+            tf.print(
+                "\n",
+                self.model.loss.name,
+                self.param_name,
+                self.model.loss.__getattribute__(self.param_name),
+            )
+        super(LossParamLog, self).on_train_batch_end(epoch, logs)
+
+    def get_config(self):
+        return {
+            "param_name": self.param_name,
+            "rate": self.rate,
+        }
diff --git a/deel/lip/initializers.py b/deel/lip/initializers.py
index 18eb80cb..09242499 100644
--- a/deel/lip/initializers.py
+++ b/deel/lip/initializers.py
@@ -2,7 +2,7 @@
 # rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
 # CRIAQ and ANITI - https://www.deel.ai/
 # =====================================================================================
-from tensorflow.keras.initializers import Initializer, Orthogonal
+from tensorflow.keras.initializers import Initializer
 from tensorflow.keras import initializers
 from .normalizers import (
     reshaped_kernel_orthogonalization,
@@ -21,7 +21,7 @@ def __init__(
         eps_bjorck=DEFAULT_EPS_BJORCK,
         beta_bjorck=DEFAULT_BETA_BJORCK,
         k_coef_lip=1.0,
-        base_initializer=Orthogonal(gain=1.0, seed=None),
+        base_initializer="orthogonal",
     ) -> None:
         """
         Initialize a kernel to be 1-lipschitz orthogonal using bjorck
diff --git a/deel/lip/layers.py b/deel/lip/layers.py
index e95498ea..0f9ce78b 100644
--- a/deel/lip/layers.py
+++ b/deel/lip/layers.py
@@ -38,9 +38,16 @@
     DEFAULT_EPS_SPECTRAL,
     reshaped_kernel_orthogonalization,
     DEFAULT_BETA_BJORCK,
+    DEFAULT_MAXITER_BJORCK,
+    DEFAULT_MAXITER_SPECTRAL,
 )
 from tensorflow.keras.utils import register_keras_serializable
 
+try:
+    from keras.utils import conv_utils  # in Keras for TF >= 2.6
+except ModuleNotFoundError:
+    from tensorflow.python.keras.utils import conv_utils  # in TF.python for TF <= 2.5
+
 
 class LipschitzLayer(abc.ABC):
     """
@@ -150,6 +157,16 @@ def vanilla_export(self):
         pass
 
 
+def _check_RKO_params(eps_spectral, eps_bjorck, beta_bjorck):
+    """Assert that the RKO hyper-parameters are supported values."""
+    if eps_spectral <= 0:
+        raise ValueError("eps_spectral has to be > 0")
+    if (eps_bjorck is not None) and (eps_bjorck <= 0.0):
+        raise ValueError("eps_bjorck must be > 0")
+    if (beta_bjorck is not None) and not (0.0 < beta_bjorck <= 0.5):
+        raise ValueError("beta_bjorck must be in ]0, 0.5]")
+
+
 @register_keras_serializable("deel-lip", "SpectralDense")
 class SpectralDense(keraslayers.Dense, LipschitzLayer, Condensable):
     def __init__(
@@ -168,6 +185,8 @@ def __init__(
         eps_spectral=DEFAULT_EPS_SPECTRAL,
         eps_bjorck=DEFAULT_EPS_BJORCK,
         beta_bjorck=DEFAULT_BETA_BJORCK,
+        maxiter_spectral=DEFAULT_MAXITER_SPECTRAL,
+        maxiter_bjorck=DEFAULT_MAXITER_BJORCK,
         **kwargs
     ):
         """
@@ -198,6 +217,8 @@ def __init__(
             eps_spectral: stopping criterion for the iterative power algorithm.
             eps_bjorck: stopping criterion Bjorck algorithm.
             beta_bjorck: beta parameter in bjorck algorithm.
+            maxiter_spectral: maximum number of iterations for the power iteration.
+            maxiter_bjorck: maximum number of iterations for bjorck algorithm.
 
         Input shape:
             N-D tensor with shape: `(batch_size, ..., input_dim)`.
@@ -226,21 +247,16 @@ def __init__(
         )
         self._kwargs = kwargs
         self.set_klip_factor(k_coef_lip)
+        _check_RKO_params(eps_spectral, eps_bjorck, beta_bjorck)
         self.eps_spectral = eps_spectral
-        self.beta_bjorck = beta_bjorck
-        if (self.beta_bjorck is not None) and (
-            not ((self.beta_bjorck <= 0.5) and (self.beta_bjorck > 0.0))
-        ):
-            raise RuntimeError("beta_bjorck must be in ]0, 0.5]")
         self.eps_bjorck = eps_bjorck
-        if (self.eps_bjorck is not None) and (not self.eps_bjorck > 0.0):
-            raise RuntimeError("eps_bjorck must be in > 0")
+        self.beta_bjorck = beta_bjorck
+        self.maxiter_bjorck = maxiter_bjorck
+        self.maxiter_spectral = maxiter_spectral
         self.u = None
         self.sig = None
         self.wbar = None
         self.built = False
-        if self.eps_spectral < 0:
-            raise RuntimeError("eps_spectral has to be > 0")
 
     def build(self, input_shape):
         super(SpectralDense, self).build(input_shape)
@@ -276,6 +292,8 @@ def call(self, x, training=True):
                 self.eps_spectral,
                 self.eps_bjorck,
                 self.beta_bjorck,
+                self.maxiter_spectral,
+                self.maxiter_bjorck,
             )
             self.wbar.assign(wbar)
             self.u.assign(u)
@@ -295,6 +313,8 @@ def get_config(self):
             "eps_spectral": self.eps_spectral,
             "eps_bjorck": self.eps_bjorck,
             "beta_bjorck": self.beta_bjorck,
+            "maxiter_spectral": self.maxiter_spectral,
+            "maxiter_bjorck": self.maxiter_bjorck,
         }
         base_config = super(SpectralDense, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
@@ -307,6 +327,8 @@ def condense(self):
             self.eps_spectral,
             self.eps_bjorck,
             self.beta_bjorck,
+            self.maxiter_spectral,
+            self.maxiter_bjorck,
         )
         self.kernel.assign(wbar)
         self.u.assign(u)
@@ -329,6 +351,32 @@ def vanilla_export(self):
         return layer
 
 
+def _compute_conv_lip_factor(kernel_size, strides, input_shape, data_format):
+    """Compute the Lipschitz factor to apply on estimated Lipschitz constant in
+    convolutional layer. This factor depends on the kernel size, the strides and the
+    input shape.
+    """
+    stride = np.prod(strides)
+    kh, kw = kernel_size[0], kernel_size[1]
+    kh_div2 = (kh - 1) / 2
+    kw_div2 = (kw - 1) / 2
+
+    if data_format == "channels_last":
+        h, w = input_shape[-3], input_shape[-2]
+    elif data_format == "channels_first":
+        h, w = input_shape[-2], input_shape[-1]
+    else:
+        raise RuntimeError("data_format not understood: " % data_format)
+
+    if stride == 1:
+        return np.sqrt(
+            (w * h)
+            / ((kh * h - kh_div2 * (kh_div2 + 1)) * (kw * w - kw_div2 * (kw_div2 + 1)))
+        )
+    else:
+        return np.sqrt(1.0 / (np.ceil(kh / strides[0]) * np.ceil(kw / strides[1])))
+
+
 @register_keras_serializable("deel-lip", "SpectralConv2D")
 class SpectralConv2D(keraslayers.Conv2D, LipschitzLayer, Condensable):
     def __init__(
@@ -352,6 +400,8 @@ def __init__(
         eps_spectral=DEFAULT_EPS_SPECTRAL,
         eps_bjorck=DEFAULT_EPS_BJORCK,
         beta_bjorck=DEFAULT_BETA_BJORCK,
+        maxiter_spectral=DEFAULT_MAXITER_SPECTRAL,
+        maxiter_bjorck=DEFAULT_MAXITER_BJORCK,
         **kwargs
     ):
         """
@@ -412,6 +462,8 @@ def __init__(
             eps_spectral: stopping criterion for the iterative power algorithm.
             eps_bjorck: stopping criterion Bjorck algorithm.
             beta_bjorck: beta parameter in bjorck algorithm.
+            maxiter_spectral: maximum number of iterations for the power iteration.
+            maxiter_bjorck: maximum number of iterations for bjorck algorithm.
 
         This documentation reuse the body of the original keras.layers.Conv2D doc.
         """
@@ -420,9 +472,9 @@ def __init__(
             or (dilation_rate == [1, 1])
             or (dilation_rate == 1)
         ):
-            raise RuntimeError("NormalizedConv does not support dilation rate")
+            raise RuntimeError("SpectralConv2D does not support dilation rate")
         if padding != "same":
-            raise RuntimeError("NormalizedConv only support padding='same'")
+            raise RuntimeError("SpectralConv2D only supports padding='same'")
         super(SpectralConv2D, self).__init__(
             filters=filters,
             kernel_size=kernel_size,
@@ -446,17 +498,12 @@ def __init__(
         self.u = None
         self.sig = None
         self.wbar = None
+        _check_RKO_params(eps_spectral, eps_bjorck, beta_bjorck)
         self.eps_spectral = eps_spectral
-        self.beta_bjorck = beta_bjorck
-        if (self.beta_bjorck is not None) and (
-            not ((self.beta_bjorck <= 0.5) and (self.beta_bjorck > 0.0))
-        ):
-            raise RuntimeError("beta_bjorck must be in ]0, 0.5]")
         self.eps_bjorck = eps_bjorck
-        if (self.eps_bjorck is not None) and (not self.eps_bjorck > 0.0):
-            raise RuntimeError("eps_bjorck must be in > 0")
-        if self.eps_spectral < 0:
-            raise RuntimeError("eps_spectral has to be > 0")
+        self.beta_bjorck = beta_bjorck
+        self.maxiter_bjorck = maxiter_bjorck
+        self.maxiter_spectral = maxiter_spectral
 
     def build(self, input_shape):
         super(SpectralConv2D, self).build(input_shape)
@@ -480,33 +527,9 @@ def build(self, input_shape):
         self.built = True
 
     def _compute_lip_coef(self, input_shape=None):
-        # According to the file lipschitz_CNN.pdf
-        stride = np.prod(self.strides)
-        k1 = self.kernel_size[0]
-        k1_div2 = (k1 - 1) / 2
-        k2 = self.kernel_size[1]
-        k2_div2 = (k2 - 1) / 2
-        if self.data_format == "channels_last":
-            h = input_shape[-3]
-            w = input_shape[-2]
-        elif self.data_format == "channels_first":
-            h = input_shape[-2]
-            w = input_shape[-1]
-        else:
-            raise RuntimeError("data_format not understood: " % self.data_format)
-        if stride == 1:
-            coefLip = np.sqrt(
-                (w * h)
-                / (
-                    (k1 * h - k1_div2 * (k1_div2 + 1))
-                    * (k2 * w - k2_div2 * (k2_div2 + 1))
-                )
-            )
-        else:
-            sn1 = self.strides[0]
-            sn2 = self.strides[1]
-            coefLip = np.sqrt(1.0 / (np.ceil(k1 / sn1) * np.ceil(k2 / sn2)))
-        return coefLip
+        return _compute_conv_lip_factor(
+            self.kernel_size, self.strides, input_shape, self.data_format
+        )
 
     def call(self, x, training=True):
         if training:
@@ -517,6 +540,8 @@ def call(self, x, training=True):
                 self.eps_spectral,
                 self.eps_bjorck,
                 self.beta_bjorck,
+                self.maxiter_spectral,
+                self.maxiter_bjorck,
             )
             self.wbar.assign(wbar)
             self.u.assign(u)
@@ -543,6 +568,8 @@ def get_config(self):
             "eps_spectral": self.eps_spectral,
             "eps_bjorck": self.eps_bjorck,
             "beta_bjorck": self.beta_bjorck,
+            "maxiter_spectral": self.maxiter_spectral,
+            "maxiter_bjorck": self.maxiter_bjorck,
         }
         base_config = super(SpectralConv2D, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
@@ -555,6 +582,8 @@ def condense(self):
             self.eps_spectral,
             self.eps_bjorck,
             self.beta_bjorck,
+            self.maxiter_spectral,
+            self.maxiter_bjorck,
         )
         self.kernel.assign(wbar)
         self.u.assign(u)
@@ -582,6 +611,305 @@ def vanilla_export(self):
         return layer
 
 
+@register_keras_serializable("deel-lip", "SpectralConv2DTranspose")
+class SpectralConv2DTranspose(keraslayers.Conv2DTranspose, LipschitzLayer, Condensable):
+    def __init__(
+        self,
+        filters,
+        kernel_size,
+        strides=(1, 1),
+        padding="same",
+        output_padding=None,
+        data_format=None,
+        dilation_rate=(1, 1),
+        activation=None,
+        use_bias=True,
+        kernel_initializer=SpectralInitializer(),
+        bias_initializer="zeros",
+        kernel_regularizer=None,
+        bias_regularizer=None,
+        activity_regularizer=None,
+        kernel_constraint=None,
+        bias_constraint=None,
+        k_coef_lip=1.0,
+        eps_spectral=DEFAULT_EPS_SPECTRAL,
+        eps_bjorck=DEFAULT_EPS_BJORCK,
+        beta_bjorck=DEFAULT_BETA_BJORCK,
+        maxiter_spectral=DEFAULT_MAXITER_SPECTRAL,
+        maxiter_bjorck=DEFAULT_MAXITER_BJORCK,
+        **kwargs
+    ):
+        """
+        This class is a Conv2DTranspose layer constrained such that all singular values
+        of its kernel are 1. The computation is based on Björck orthogonalization
+        algorithm.
+
+        The computation is done in three steps:
+        1. reduce the largest singular value to 1, using iterated power method.
+        2. increase other singular values to 1, using Björck algorithm.
+        3. divide the output by the Lipschitz target K to ensure K-Lipschitzity.
+
+        This documentation reuses the body of the original
+        `tf.keras.layers.Conv2DTranspose` doc.
+
+        Args:
+            filters: Integer, the dimensionality of the output space
+                (i.e. the number of output filters in the convolution).
+            kernel_size: An integer or tuple/list of 2 integers, specifying the
+                height and width of the 2D convolution window.
+                Can be a single integer to specify the same value for
+                all spatial dimensions.
+            strides: An integer or tuple/list of 2 integers,
+                specifying the strides of the convolution along the height and width.
+                Can be a single integer to specify the same value for
+                all spatial dimensions.
+            padding: only `"same"` padding is supported in this Lipschitz layer
+                (case-insensitive).
+            output_padding: if set to `None` (default), the output shape is inferred.
+                Only `None` value is supported in this Lipschitz layer.
+            data_format: A string,
+                one of `channels_last` (default) or `channels_first`.
+                The ordering of the dimensions in the inputs.
+                `channels_last` corresponds to inputs with shape
+                `(batch, height, width, channels)` while `channels_first`
+                corresponds to inputs with shape
+                `(batch, channels, height, width)`.
+                It defaults to the `image_data_format` value found in your
+                Keras config file at `~/.keras/keras.json`.
+                If you never set it, then it will be "channels_last".
+            dilation_rate: an integer, specifying the dilation rate for all spatial
+                dimensions for dilated convolution. This Lipschitz layer does not
+                support dilation rate != 1.
+            activation: Activation function to use.
+                If you don't specify anything, no activation is applied
+                (see `keras.activations`).
+            use_bias: Boolean, whether the layer uses a bias vector.
+            kernel_initializer: Initializer for the `kernel` weights matrix
+                (see `keras.initializers`). Defaults to `SpectralInitializer`.
+            bias_initializer: Initializer for the bias vector
+                (see `keras.initializers`). Defaults to 'zeros'.
+            kernel_regularizer: Regularizer function applied to
+                the `kernel` weights matrix (see `keras.regularizers`).
+            bias_regularizer: Regularizer function applied to the bias vector
+                (see `keras.regularizers`).
+            activity_regularizer: Regularizer function applied to
+                the output of the layer (its "activation") (see `keras.regularizers`).
+            kernel_constraint: Constraint function applied to the kernel matrix
+                (see `keras.constraints`).
+            bias_constraint: Constraint function applied to the bias vector
+                (see `keras.constraints`).
+            k_coef_lip: Lipschitz constant to ensure
+            eps_spectral: stopping criterion for the iterative power algorithm.
+            eps_bjorck: stopping criterion Björck algorithm.
+            beta_bjorck: beta parameter in Björck algorithm.
+            maxiter_spectral: maximum number of iterations for the power iteration.
+            maxiter_bjorck: maximum number of iterations for bjorck algorithm.
+        """
+        super().__init__(
+            filters,
+            kernel_size,
+            strides,
+            padding,
+            output_padding,
+            data_format,
+            dilation_rate,
+            activation,
+            use_bias,
+            kernel_initializer,
+            bias_initializer,
+            kernel_regularizer,
+            bias_regularizer,
+            activity_regularizer,
+            kernel_constraint,
+            bias_constraint,
+            **kwargs
+        )
+
+        if self.dilation_rate != (1, 1):
+            raise ValueError("SpectralConv2DTranspose does not support dilation rate")
+        if self.padding != "same":
+            raise ValueError("SpectralConv2DTranspose only supports padding='same'")
+        if self.output_padding is not None:
+            raise ValueError(
+                "SpectralConv2DTranspose only supports output_padding=None"
+            )
+        self.set_klip_factor(k_coef_lip)
+        self.u = None
+        self.sig = None
+        self.wbar = None
+        _check_RKO_params(eps_spectral, eps_bjorck, beta_bjorck)
+        self.eps_spectral = eps_spectral
+        self.eps_bjorck = eps_bjorck
+        self.beta_bjorck = beta_bjorck
+        self.maxiter_bjorck = maxiter_bjorck
+        self.maxiter_spectral = maxiter_spectral
+        self._kwargs = kwargs
+
+    def build(self, input_shape):
+        super().build(input_shape)
+        self._init_lip_coef(input_shape)
+        self.u = self.add_weight(
+            shape=tuple([1, self.kernel.shape.as_list()[2]]),
+            initializer=RandomNormal(0, 1),
+            name="sn",
+            trainable=False,
+            dtype=self.dtype,
+        )
+
+        self.sig = self.add_weight(
+            shape=tuple([1, 1]),  # maximum spectral  value
+            name="sigma",
+            trainable=False,
+            dtype=self.dtype,
+        )
+        self.sig.assign([[1.0]])
+        self.wbar = tf.Variable(self.kernel.read_value(), trainable=False)
+
+    def _compute_lip_coef(self, input_shape=None):
+        return _compute_conv_lip_factor(
+            self.kernel_size, self.strides, input_shape, self.data_format
+        )
+
+    def call(self, inputs, training=True):
+        if training:
+            kernel_reshaped = tf.transpose(self.kernel, [0, 1, 3, 2])
+            wbar, u, sigma = reshaped_kernel_orthogonalization(
+                kernel_reshaped,
+                self.u,
+                self._get_coef(),
+                self.eps_spectral,
+                self.eps_bjorck,
+                self.beta_bjorck,
+                self.maxiter_spectral,
+                self.maxiter_bjorck,
+            )
+            wbar = tf.transpose(wbar, [0, 1, 3, 2])
+            self.wbar.assign(wbar)
+            self.u.assign(u)
+            self.sig.assign(sigma)
+        else:
+            wbar = self.wbar
+
+        # Apply conv2D_transpose operation on constrained weights
+        # (code from TF/Keras 2.9.1)
+        inputs_shape = tf.shape(inputs)
+        batch_size = inputs_shape[0]
+        if self.data_format == "channels_first":
+            h_axis, w_axis = 2, 3
+        else:
+            h_axis, w_axis = 1, 2
+
+        height, width = None, None
+        if inputs.shape.rank is not None:
+            dims = inputs.shape.as_list()
+            height = dims[h_axis]
+            width = dims[w_axis]
+        height = height if height is not None else inputs_shape[h_axis]
+        width = width if width is not None else inputs_shape[w_axis]
+
+        kernel_h, kernel_w = self.kernel_size
+        stride_h, stride_w = self.strides
+
+        if self.output_padding is None:
+            out_pad_h = out_pad_w = None
+        else:
+            out_pad_h, out_pad_w = self.output_padding
+
+        # Infer the dynamic output shape:
+        out_height = conv_utils.deconv_output_length(
+            height,
+            kernel_h,
+            padding=self.padding,
+            output_padding=out_pad_h,
+            stride=stride_h,
+            dilation=self.dilation_rate[0],
+        )
+        out_width = conv_utils.deconv_output_length(
+            width,
+            kernel_w,
+            padding=self.padding,
+            output_padding=out_pad_w,
+            stride=stride_w,
+            dilation=self.dilation_rate[1],
+        )
+        if self.data_format == "channels_first":
+            output_shape = (batch_size, self.filters, out_height, out_width)
+        else:
+            output_shape = (batch_size, out_height, out_width, self.filters)
+        output_shape_tensor = tf.stack(output_shape)
+
+        outputs = K.conv2d_transpose(
+            inputs,
+            wbar,
+            output_shape_tensor,
+            strides=self.strides,
+            padding=self.padding,
+            data_format=self.data_format,
+            dilation_rate=self.dilation_rate,
+        )
+
+        if not tf.executing_eagerly():
+            # Infer the static output shape:
+            out_shape = self.compute_output_shape(inputs.shape)
+            outputs.set_shape(out_shape)
+
+        if self.use_bias:
+            outputs = tf.nn.bias_add(
+                outputs,
+                self.bias,
+                data_format=conv_utils.convert_data_format(self.data_format, ndim=4),
+            )
+
+        if self.activation is not None:
+            return self.activation(outputs)
+        return outputs
+
+    def get_config(self):
+        config = {
+            "k_coef_lip": self.k_coef_lip,
+            "eps_spectral": self.eps_spectral,
+            "eps_bjorck": self.eps_bjorck,
+            "beta_bjorck": self.beta_bjorck,
+            "maxiter_spectral": self.maxiter_spectral,
+            "maxiter_bjorck": self.maxiter_bjorck,
+        }
+        base_config = super().get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+    def condense(self):
+        wbar, u, sigma = reshaped_kernel_orthogonalization(
+            self.kernel,
+            self.u,
+            self._get_coef(),
+            self.eps_spectral,
+            self.eps_bjorck,
+            self.beta_bjorck,
+            self.maxiter_spectral,
+            self.maxiter_bjorck,
+        )
+        self.kernel.assign(wbar)
+        self.u.assign(u)
+        self.sig.assign(sigma)
+
+    def vanilla_export(self):
+        self._kwargs["name"] = self.name
+        layer = keraslayers.Conv2DTranspose(
+            filters=self.filters,
+            kernel_size=self.kernel_size,
+            strides=self.strides,
+            padding=self.padding,
+            data_format=self.data_format,
+            activation=self.activation,
+            use_bias=self.use_bias,
+            **self._kwargs
+        )
+        layer.build(self.input_shape)
+        layer.kernel.assign(self.wbar)
+        if layer.use_bias:
+            layer.bias.assign(self.bias)
+        return layer
+
+
 @register_keras_serializable("deel-lip", "FrobeniusDense")
 class FrobeniusDense(keraslayers.Dense, LipschitzLayer, Condensable):
     """
@@ -716,15 +1044,15 @@ def __init__(
         **kwargs
     ):
         if not ((strides == (1, 1)) or (strides == [1, 1]) or (strides == 1)):
-            raise RuntimeError("NormalizedConv does not support strides")
+            raise RuntimeError("FrobeniusConv2D does not support strides")
         if not (
             (dilation_rate == (1, 1))
             or (dilation_rate == [1, 1])
             or (dilation_rate == 1)
         ):
-            raise RuntimeError("NormalizedConv does not support dilation rate")
+            raise RuntimeError("FrobeniusConv2D does not support dilation rate")
         if padding != "same":
-            raise RuntimeError("NormalizedConv only support padding='same'")
+            raise RuntimeError("FrobeniusConv2D only supports padding='same'")
         if not (
             (kernel_constraint is None)
             or isinstance(kernel_constraint, SpectralConstraint)
@@ -762,7 +1090,9 @@ def build(self, input_shape):
         self.built = True
 
     def _compute_lip_coef(self, input_shape=None):
-        return SpectralConv2D._compute_lip_coef(self, input_shape)
+        return _compute_conv_lip_factor(
+            self.kernel_size, self.strides, input_shape, self.data_format
+        )
 
     def call(self, x, training=True):
         if training:
@@ -855,7 +1185,7 @@ def __init__(
         if not ((strides == pool_size) or (strides is None)):
             raise RuntimeError("stride must be equal to pool_size")
         if padding != "valid":
-            raise RuntimeError("ScaledAveragePooling2D only support padding='valid'")
+            raise RuntimeError("ScaledAveragePooling2D only supports padding='valid'")
         super(ScaledAveragePooling2D, self).__init__(
             pool_size=pool_size,
             strides=pool_size,
@@ -943,7 +1273,7 @@ def __init__(
         if not ((strides == pool_size) or (strides is None)):
             raise RuntimeError("stride must be equal to pool_size")
         if padding != "valid":
-            raise RuntimeError("NormalizedConv only support padding='valid'")
+            raise RuntimeError("ScaledL2NormPooling2D only supports padding='valid'")
         if eps_grad_sqrt < 0.0:
             raise RuntimeError("eps_grad_sqrt must be positive")
         super(ScaledL2NormPooling2D, self).__init__(
diff --git a/deel/lip/losses.py b/deel/lip/losses.py
index 836ac9f4..caa88766 100644
--- a/deel/lip/losses.py
+++ b/deel/lip/losses.py
@@ -9,8 +9,7 @@
 from functools import partial
 import numpy as np
 import tensorflow as tf
-from tensorflow.keras.losses import Loss
-from tensorflow.keras.losses import Reduction
+from tensorflow.keras.losses import categorical_crossentropy, Loss, Reduction
 from tensorflow.keras.utils import register_keras_serializable
 
 
@@ -145,34 +144,58 @@ def __init__(
             name: passed to tf.keras.Loss constructor
 
         """
-        self.alpha = alpha
-        self.min_margin = min_margin
+        self.alpha = tf.Variable(alpha, dtype=tf.float32)
+        self.min_margin = tf.Variable(min_margin, dtype=tf.float32)
         self.multi_gpu = multi_gpu
         self.KRloss = KR(multi_gpu=multi_gpu)
-        self.hingeloss = HingeMargin(min_margin)
+        if alpha == np.inf:  # alpha = inf => hinge only
+            self.fct = partial(hinge_margin, min_margin=self.min_margin)
+        else:
+            self.fct = self.hkr
         super(HKR, self).__init__(reduction=reduction, name=name)
 
     @tf.function
+    def hkr(self, y_true, y_pred):
+        a = -self.KRloss.call(y_true, y_pred)
+        b = hinge_margin(y_true, y_pred, self.min_margin)
+        return a + self.alpha * b
+
     def call(self, y_true, y_pred):
-        if self.alpha == np.inf:  # alpha = inf => hinge only
-            return self.hingeloss.call(y_true, y_pred)
-        elif self.alpha == 0.0:  # alpha = 0 => KR only
-            return -self.KRloss.call(y_true, y_pred)
-        else:
-            kr = -self.KRloss.call(y_true, y_pred)
-            hinge = self.hingeloss.call(y_true, y_pred)
-            return kr + self.alpha * hinge
+        return self.fct(y_true, y_pred)
 
     def get_config(self):
         config = {
-            "alpha": self.alpha,
-            "min_margin": self.min_margin,
+            "alpha": self.alpha.numpy(),
+            "min_margin": self.min_margin.numpy(),
             "multi_gpu": self.multi_gpu,
         }
         base_config = super(HKR, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
 
 
+def hinge_margin(y_true, y_pred, min_margin):
+    """Compute the element-wise binary hinge margin loss.
+
+    Note that `y_true` and `y_pred` must be of rank 2: (batch_size, 1) or
+    (batch_size, C) for multilabel classification (with C categories).
+    `y_true` accepts label values in (0, 1), (-1, 1), or pre-processed with the
+    :func:`deel.lip.utils.process_labels_for_multi_gpu()` function.
+
+    Args:
+        min_margin: positive float, margin to enforce.
+
+    Returns:
+        Element-wise hinge margin loss value.
+
+    """
+    sign = tf.where(y_true > 0, 1, -1)
+    sign = tf.cast(sign, y_pred.dtype)
+    hinge = tf.nn.relu(min_margin / 2.0 - sign * y_pred)
+    # In binary case (`y_true` of shape (batch_size, 1)), `tf.reduce_mean(axis=-1)`
+    # behaves like `tf.squeeze` to return element-wise loss of shape (batch_size, ).
+    return tf.reduce_mean(hinge, axis=-1)
+
+
 @register_keras_serializable("deel-lip", "HingeMargin")
 class HingeMargin(Loss):
     def __init__(self, min_margin=1.0, reduction=Reduction.AUTO, name="HingeMargin"):
@@ -194,21 +217,16 @@ def __init__(self, min_margin=1.0, reduction=Reduction.AUTO, name="HingeMargin")
             name: passed to tf.keras.Loss constructor
 
         """
-        self.min_margin = min_margin
+        self.min_margin = tf.Variable(min_margin, dtype=tf.float32)
         super(HingeMargin, self).__init__(reduction=reduction, name=name)
 
     @tf.function
     def call(self, y_true, y_pred):
-        sign = tf.where(y_true > 0, 1, -1)
-        sign = tf.cast(sign, y_pred.dtype)
-        hinge = tf.nn.relu(self.min_margin - sign * y_pred)
-        # In binary case (`y_true` of shape (batch_size, 1)), `tf.reduce_mean(axis=-1)`
-        # behaves like `tf.squeeze` to return element-wise loss of shape (batch_size, ).
-        return tf.reduce_mean(hinge, axis=-1)
+        return hinge_margin(y_true, y_pred, self.min_margin)
 
     def get_config(self):
         config = {
-            "min_margin": self.min_margin,
+            "min_margin": self.min_margin.numpy(),
         }
         base_config = super(HingeMargin, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
@@ -253,6 +271,31 @@ def get_config(self):
         return dict(list(base_config.items()) + list(config.items()))
 
 
+def multiclass_hinge(y_true, y_pred, min_margin):
+    """Compute the multi-class hinge margin loss.
+
+    `y_true` and `y_pred` must be of shape (batch_size, # classes).
+    Note that `y_true` should be one-hot encoded or pre-processed with the
+    :func:`deel.lip.utils.process_labels_for_multi_gpu()` function.
+
+    Args:
+        y_true: tensor of true targets of shape (batch_size, # classes)
+        y_pred: tensor of predicted targets of shape (batch_size, # classes)
+        min_margin: positive float, margin to enforce.
+
+    Returns:
+        Element-wise multi-class hinge margin loss value.
+    """
+    sign = tf.where(y_true > 0, 1, -1)
+    sign = tf.cast(sign, y_pred.dtype)
+    # compute the elementwise hinge term
+    hinge = tf.nn.relu(min_margin / 2.0 - sign * y_pred)
+    # reweight positive elements
+    factor = y_pred.shape[-1] - 1.0
+    hinge = tf.where(sign > 0, hinge * factor, hinge)
+    return tf.reduce_mean(hinge, axis=-1)
+
+
 @register_keras_serializable("deel-lip", "MulticlassHinge")
 class MulticlassHinge(Loss):
     def __init__(
@@ -274,26 +317,16 @@ def __init__(
             name: passed to tf.keras.Loss constructor
 
         """
-        self.min_margin = min_margin
+        self.min_margin = tf.Variable(min_margin, dtype=tf.float32)
         super(MulticlassHinge, self).__init__(reduction=reduction, name=name)
 
     @tf.function
     def call(self, y_true, y_pred):
-        sign = tf.where(y_true > 0, 1, -1)
-        sign = tf.cast(sign, y_pred.dtype)
-        # compute the elementwise hinge term
-        hinge = tf.nn.relu(self.min_margin - sign * y_pred)
-        # reweight positive elements
-        if (len(tf.shape(y_pred)) == 2) and (tf.shape(y_pred)[-1] != 1):
-            factor = y_true.shape[-1] - 1.0
-        else:
-            factor = 1.0
-        hinge = tf.where(sign > 0, hinge * factor, hinge)
-        return tf.reduce_mean(hinge, axis=-1)
+        return multiclass_hinge(y_true, y_pred, self.min_margin)
 
     def get_config(self):
         config = {
-            "min_margin": self.min_margin,
+            "min_margin": self.min_margin.numpy(),
         }
         base_config = super(MulticlassHinge, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
@@ -328,28 +361,29 @@ class and averaging the results.
             name: passed to tf.keras.Loss constructor
 
         """
-        self.alpha = alpha
-        self.min_margin = min_margin
+        self.alpha = tf.Variable(alpha, dtype=tf.float32)
+        self.min_margin = tf.Variable(min_margin, dtype=tf.float32)
         self.multi_gpu = multi_gpu
-        self.hingeloss = MulticlassHinge(self.min_margin)
-        self.KRloss = MulticlassKR(multi_gpu=multi_gpu, name=name)
+        self.KRloss = MulticlassKR(multi_gpu=multi_gpu, reduction=reduction, name=name)
+        if alpha == np.inf:  # alpha = inf => hinge only
+            self.fct = partial(multiclass_hinge, min_margin=self.min_margin)
+        else:
+            self.fct = self.hkr
         super(MulticlassHKR, self).__init__(reduction=reduction, name=name)
 
     @tf.function
+    def hkr(self, y_true, y_pred):
+        a = -self.KRloss.call(y_true, y_pred)
+        b = multiclass_hinge(y_true, y_pred, self.min_margin)
+        return a + self.alpha * b
+
     def call(self, y_true, y_pred):
-        if self.alpha == np.inf:  # alpha = inf => hinge only
-            return self.hingeloss.call(y_true, y_pred)
-        elif self.alpha == 0.0:  # alpha = 0 => KR only
-            return -self.KRloss.call(y_true, y_pred)
-        else:
-            kr = -self.KRloss.call(y_true, y_pred)
-            hinge = self.hingeloss.call(y_true, y_pred)
-            return kr + self.alpha * hinge
+        return self.fct(y_true, y_pred)
 
     def get_config(self):
         config = {
-            "alpha": self.alpha,
-            "min_margin": self.min_margin,
+            "alpha": self.alpha.numpy(),
+            "min_margin": self.min_margin.numpy(),
             "multi_gpu": self.multi_gpu,
         }
         base_config = super(MulticlassHKR, self).get_config()
@@ -372,7 +406,7 @@ def __init__(self, min_margin=1.0, reduction=Reduction.AUTO, name="MultiMargin")
             name: passed to tf.keras.Loss constructor
 
         """
-        self.min_margin = min_margin
+        self.min_margin = tf.Variable(min_margin, dtype=tf.float32)
         super(MultiMargin, self).__init__(reduction=reduction, name=name)
 
     @tf.function
@@ -390,7 +424,70 @@ def call(self, y_true, y_pred):
 
     def get_config(self):
         config = {
-            "min_margin": self.min_margin,
+            "min_margin": self.min_margin.numpy(),
         }
         base_config = super(MultiMargin, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
+
+
+@register_keras_serializable("deel-lip", "CategoricalHinge")
+class CategoricalHinge(Loss):
+    def __init__(self, min_margin, reduction=Reduction.AUTO, name="CategoricalHinge"):
+        """
+        Similar to original categorical hinge, but with a settable margin parameter.
+        This implementation is sligthly different from the Keras one.
+
+        `y_true` and `y_pred` must be of shape (batch_size, # classes).
+        Note that `y_true` should be one-hot encoded or pre-processed with the
+        :func:`deel.lip.utils.process_labels_for_multi_gpu()` function.
+
+        Args:
+            min_margin: positive float, margin parameter.
+            reduction: reduction of the loss, passed to original loss.
+            name: name of the loss
+        """
+        self.min_margin = tf.Variable(min_margin, dtype=tf.float32)
+        super(CategoricalHinge, self).__init__(name=name, reduction=reduction)
+
+    def call(self, y_true, y_pred):
+        mask = tf.where(y_true > 0, 1, 0)
+        mask = tf.cast(mask, y_pred.dtype)
+        pos = tf.reduce_sum(mask * y_pred, axis=-1)
+        neg = tf.reduce_max(tf.where(mask > 0, tf.float32.min, y_pred), axis=-1)
+        return tf.nn.relu(self.min_margin - (pos - neg))
+
+    def get_config(self):
+        config = {"min_margin": self.min_margin.numpy()}
+        base_config = super(CategoricalHinge, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+@register_keras_serializable("deel-lip", "TauCategoricalCrossentropy")
+class TauCategoricalCrossentropy(Loss):
+    def __init__(
+        self, tau, reduction=Reduction.AUTO, name="TauCategoricalCrossentropy"
+    ):
+        """
+        Similar to original categorical crossentropy, but with a settable temperature
+        parameter.
+
+        Args:
+            tau: temperature parameter.
+            reduction: reduction of the loss, passed to original loss.
+            name: name of the loss
+        """
+        self.tau = tf.Variable(tau, dtype=tf.float32)
+        super(TauCategoricalCrossentropy, self).__init__(name=name, reduction=reduction)
+
+    def call(self, y_true, y_pred, *args, **kwargs):
+        return (
+            categorical_crossentropy(
+                y_true, self.tau * y_pred, from_logits=True, *args, **kwargs
+            )
+            / self.tau
+        )
+
+    def get_config(self):
+        config = {"tau": self.tau.numpy()}
+        base_config = super(TauCategoricalCrossentropy, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
diff --git a/deel/lip/normalizers.py b/deel/lip/normalizers.py
index 9aed3242..6d6a11f8 100644
--- a/deel/lip/normalizers.py
+++ b/deel/lip/normalizers.py
@@ -12,6 +12,42 @@
 DEFAULT_BETA_BJORCK = 0.5
 DEFAULT_EPS_SPECTRAL = 1e-3
 DEFAULT_EPS_BJORCK = 1e-3
+DEFAULT_MAXITER_BJORCK = 15
+DEFAULT_MAXITER_SPECTRAL = 10
+SWAP_MEMORY = True
+STOP_GRAD_SPECTRAL = True
+
+
+def set_swap_memory(value: bool):
+    """
+    Set the global SWAP_MEMORY to values. This function must be called before
+    constructing the model (first call of `reshaped_kernel_orthogonalization`) in
+    order to be accounted.
+
+    Args:
+        value: boolean that will be used as the swap_memory parameter in while loops
+            in spectral and bjorck algorithms.
+
+    """
+    global SWAP_MEMORY
+    SWAP_MEMORY = value
+
+
+def set_stop_grad_spectral(value: bool):
+    """
+    Set the global STOP_GRAD_SPECTRAL to values. This function must be called before
+    constructing the model (first call of `reshaped_kernel_orthogonalization`) in
+    order to be accounted.
+
+    Args:
+        value: boolean, when set to True, disable back-propagation through the power
+            iteration algorithm. The back-propagation will account how updates affects
+            the maximum singular value but not how it affects the largest singular
+            vector. When set to False, back-propagate through the while loop.
+
+    """
+    global STOP_GRAD_SPECTRAL
+    STOP_GRAD_SPECTRAL = value
 
 
 def reshaped_kernel_orthogonalization(
@@ -21,6 +57,8 @@ def reshaped_kernel_orthogonalization(
     eps_spectral=DEFAULT_EPS_SPECTRAL,
     eps_bjorck=DEFAULT_EPS_BJORCK,
     beta=DEFAULT_BETA_BJORCK,
+    maxiter_spectral=DEFAULT_MAXITER_SPECTRAL,
+    maxiter_bjorck=DEFAULT_MAXITER_BJORCK,
 ):
     """
     Perform reshaped kernel orthogonalization (RKO) to the kernel given as input. It
@@ -35,20 +73,38 @@ def reshaped_kernel_orthogonalization(
         eps_spectral: stopping criterion in spectral algorithm
         eps_bjorck: stopping criterion in bjorck algorithm
         beta: the beta used in the bjorck algorithm
+        maxiter_spectral: maximum number of iterations for the power iteration
+        maxiter_bjorck: maximum number of iterations for bjorck algorithm
 
     Returns: the orthogonalized kernel, the new u, and sigma which is the largest
         singular value
 
     """
-    W_bar, u, sigma = spectral_normalization(kernel, u, eps=eps_spectral)
+    W_shape = kernel.shape
+    # Flatten the Tensor
+    W_reshaped = tf.reshape(kernel, [-1, W_shape[-1]])
+    W_bar, u, sigma = spectral_normalization(
+        W_reshaped, u, eps=eps_spectral, maxiter=maxiter_spectral
+    )
     if (eps_bjorck is not None) and (beta is not None):
-        W_bar = bjorck_normalization(W_bar, eps=eps_bjorck, beta=beta)
+        W_bar = bjorck_normalization(
+            W_bar, eps=eps_bjorck, beta=beta, maxiter=maxiter_bjorck
+        )
     W_bar = W_bar * adjustment_coef
     W_bar = K.reshape(W_bar, kernel.shape)
     return W_bar, u, sigma
 
 
-def bjorck_normalization(w, eps=DEFAULT_EPS_BJORCK, beta=DEFAULT_BETA_BJORCK):
+def _wwtw(w):
+    if w.shape[0] > w.shape[1]:
+        return w @ (tf.transpose(w) @ w)
+    else:
+        return (w @ tf.transpose(w)) @ w
+
+
+def bjorck_normalization(
+    w, eps=DEFAULT_EPS_BJORCK, beta=DEFAULT_BETA_BJORCK, maxiter=DEFAULT_MAXITER_BJORCK
+):
     """
     apply Bjorck normalization on w.
 
@@ -57,6 +113,7 @@ def bjorck_normalization(w, eps=DEFAULT_EPS_BJORCK, beta=DEFAULT_BETA_BJORCK):
             max_eigenval(w) ~= 1
         eps: epsilon stopping criterion: norm(wt - wt-1) must be less than eps
         beta: beta used in each iteration, must be in the interval ]0, 0.5]
+        maxiter: maximum number of iterations for the algorithm
 
     Returns:
         the orthonormal weights
@@ -73,17 +130,22 @@ def cond(w, old_w):
     # define the loop body
     def body(w, old_w):
         old_w = w
-        w = (1 + beta) * w - beta * w @ tf.transpose(w) @ w
+        w = (1 + beta) * w - beta * _wwtw(w)
         return w, old_w
 
     # apply the loop
     w, old_w = tf.while_loop(
-        cond, body, (w, old_w), parallel_iterations=1, maximum_iterations=30
+        cond,
+        body,
+        (w, old_w),
+        parallel_iterations=30,
+        maximum_iterations=maxiter,
+        swap_memory=SWAP_MEMORY,
     )
     return w
 
 
-def _power_iteration(w, u, eps=DEFAULT_EPS_SPECTRAL):
+def _power_iteration(w, u, eps=DEFAULT_EPS_SPECTRAL, maxiter=DEFAULT_MAXITER_SPECTRAL):
     """
     Internal function that performs the power iteration algorithm.
 
@@ -91,6 +153,7 @@ def _power_iteration(w, u, eps=DEFAULT_EPS_SPECTRAL):
         w: weights matrix that we want to find eigen vector
         u: initialization of the eigen vector
         eps: epsilon stopping criterion: norm(ut - ut-1) must be less than eps
+        maxiter: maximum number of iterations for the algorithm
 
     Returns:
          u and v corresponding to the maximum eigenvalue
@@ -98,8 +161,14 @@ def _power_iteration(w, u, eps=DEFAULT_EPS_SPECTRAL):
     """
     # build _u and _v (_v is size of _u@tf.transpose(w), will be set on the first body
     # iteration)
+    if u is None:
+        u = tf.linalg.l2_normalize(
+            tf.random.uniform(
+                shape=(1, w.shape[-1]), minval=0.0, maxval=1.0, dtype=w.dtype
+            )
+        )
     _u = u
-    _v = tf.zeros(u.shape[:-1] + (w.shape[0],), dtype=w.dtype)
+    _v = tf.zeros((1,) + (w.shape[0],), dtype=w.dtype)
 
     # create a fake old_w that doesn't pass the loop condition
     # it won't affect computation as the first action done in the loop overwrite it.
@@ -118,38 +187,44 @@ def body(_u, _v, _old_u):
 
     # apply the loop
     _u, _v, _old_u = tf.while_loop(
-        cond, body, (_u, _v, _old_u), parallel_iterations=1, maximum_iterations=30
+        cond,
+        body,
+        (_u, _v, _old_u),
+        parallel_iterations=30,
+        maximum_iterations=maxiter,
+        swap_memory=SWAP_MEMORY,
     )
+    if STOP_GRAD_SPECTRAL:
+        _u = tf.stop_gradient(_u)
+        _v = tf.stop_gradient(_v)
     return _u, _v
 
 
-def spectral_normalization(kernel, u, eps=DEFAULT_EPS_SPECTRAL):
+def spectral_normalization(
+    kernel, u, eps=DEFAULT_EPS_SPECTRAL, maxiter=DEFAULT_MAXITER_SPECTRAL
+):
     """
     Normalize the kernel to have it's max eigenvalue == 1.
 
     Args:
-        kernel: the kernel to normalize
+        kernel: the kernel to normalize, assuming a 2D kernel
         u: initialization for the max eigen vector
         eps: epsilon stopping criterion: norm(ut - ut-1) must be less than eps
+        maxiter: maximum number of iterations for the algorithm
 
     Returns:
         the normalized kernel w_bar, it's shape, the maximum eigen vector, and the
         maximum singular value
 
     """
-    W_shape = kernel.shape
-    if u is None:
-        u = tf.linalg.l2_normalize(tf.ones(shape=tuple([1, W_shape[-1]])))
-    # Flatten the Tensor
-    W_reshaped = tf.reshape(kernel, [-1, W_shape[-1]])
-    _u, _v = _power_iteration(W_reshaped, u, eps)
+    _u, _v = _power_iteration(kernel, u, eps, maxiter)
     # compute Sigma
-    sigma = _v @ W_reshaped
+    sigma = _v @ kernel
     sigma = sigma @ tf.transpose(_u)
     # normalize it
     # we assume that in the worst case we converged to sigma + eps (as u and v are
     # normalized after each iteration)
     # in order to be sure that operator norm of W_bar is strictly less than one we
     # use sigma + eps, which ensure stability of the bjorck even when beta=0.5
-    W_bar = W_reshaped / (sigma + eps)
+    W_bar = kernel / (sigma + eps)
     return W_bar, _u, sigma
diff --git a/deel/lip/regularizers.py b/deel/lip/regularizers.py
new file mode 100644
index 00000000..a72dc079
--- /dev/null
+++ b/deel/lip/regularizers.py
@@ -0,0 +1,254 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+# =====================================================================================
+
+import warnings
+from abc import ABC, abstractmethod
+
+import tensorflow as tf
+from tensorflow.keras.regularizers import Regularizer
+from tensorflow.keras.utils import register_keras_serializable
+
+
+class Lorth(ABC):
+    def __init__(self, dim, kernel_shape=None, stride=1, conv_transpose=False) -> None:
+        """
+        Base class for Lorth regularization. Not meant to be used standalone.
+
+        Ref. Achour & al., Existence, Stability And Scalability Of Orthogonal
+        Convolutional Neural Networks (2022).
+        https://www.jmlr.org/papers/v23/22-0026.html
+
+        Args:
+            dim (int): the rank of the convolution, e.g. "2" for 2D convolution.
+            kernel_shape: the shape of the kernel.
+            stride (int): stride used in the associated convolution
+            conv_transpose (bool): whether the kernel is from a transposed convolution.
+        """
+        super(Lorth, self).__init__()
+        self.dim = dim
+        self.stride = stride
+        self.conv_transpose = conv_transpose
+        self.set_kernel_shape(kernel_shape)
+
+    def _get_kernel_shape(self):
+        """Return the kernel size, the number of input channels and output channels"""
+        return [self.kernel_shape[i] for i in (0, -2, -1)]
+
+    def _compute_delta(self):
+        """delta is positive in CO case, zero in RO case."""
+        _, C, M = self._get_kernel_shape()
+        if not self.conv_transpose:
+            delta = M - (self.stride**self.dim) * C
+        else:
+            delta = C - (self.stride**self.dim) * M
+        return max(0, delta)
+
+    def _alphaNormSpectral(self):
+        """delta is positive in CO case, zero in RO case."""
+        R, C1, M1 = self._get_kernel_shape()
+        if not self.conv_transpose:
+            C, M = C1, M1
+        else:
+            C, M = M1, C1
+        alpha = (
+            R * C * M
+        )  # huge value to get the minimum value in case of square matrix
+        if M - (self.stride**self.dim) * C <= 0:  # RO case
+            alpha = M * (2 * ((R - 1) // self.stride) + 1) ** self.dim
+        if M - (self.stride**self.dim) * C >= 0:  # CO case
+            alpha = min(alpha, C * (2 * R - 1) ** self.dim)
+        return alpha
+
+    def _check_if_orthconv_exists(self):
+        """check the existence of Orthogonal convolution (for circular padding)"""
+        R, C, M = self._get_kernel_shape()
+        msg = "Impossible {} configuration for orthogonal convolution."
+        if C * self.stride**self.dim >= M:  # RO case
+            if M > C * (R**self.dim):
+                raise RuntimeError(msg.format("RO"))
+        else:  # CO case
+            if self.stride > R:
+                raise RuntimeError(msg.format("CO"))
+        if C * (self.stride**self.dim) == M:  # square case
+            warnings.warn(
+                "LorthRegularizer: warning configuration C*S^2=M is hard to optimize."
+            )
+
+    def set_kernel_shape(self, shape):
+        if shape is None:
+            self.kernel_shape, self.padding, self.delta = None, None, None
+            return
+
+        R = shape[0]
+        self.kernel_shape = shape
+        self.padding = ((R - 1) // self.stride) * self.stride
+        self.delta = self._compute_delta()
+        self.alphaNormSpectral = self._alphaNormSpectral()
+
+        # Assertions on kernel shape and existence of orthogonal convolution
+        assert R & 1, "Lorth regularizer requires odd kernels. Receives " + str(R)
+        self._check_if_orthconv_exists()
+
+    @abstractmethod
+    def _compute_conv_kk(self, w):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _compute_target(self, w, output_shape):
+        raise NotImplementedError()
+
+    def compute_lorth(self, w):
+        output = self._compute_conv_kk(w)
+        target = self._compute_target(w, output.shape)
+        return tf.reduce_sum(tf.square(output - target)) - self.delta
+
+
+class Lorth2D(Lorth):
+    def __init__(self, kernel_shape=None, stride=1, conv_transpose=False) -> None:
+        """
+        Lorth computation for 2D convolutions. Although this class allows to compute
+        the regularization term, it cannot be used as it is in a layer.
+
+        Ref. Wang & al., Orthogonal Convolutional Neural Networks (2020).
+        http://arxiv.org/abs/1911.12207
+
+        Args:
+            kernel_shape: the shape of the kernel.
+            stride (int): stride used in the associated convolution
+            conv_transpose (bool): whether the kernel is from a transposed convolution.
+        """
+        dim = 2
+        super(Lorth2D, self).__init__(dim, kernel_shape, stride, conv_transpose)
+
+    def _compute_conv_kk(self, w):
+        w_reshape = tf.transpose(w, perm=[3, 0, 1, 2])
+        w_padded = tf.pad(
+            w_reshape,
+            paddings=[
+                [0, 0],
+                [self.padding, self.padding],
+                [self.padding, self.padding],
+                [0, 0],
+            ],
+        )
+        return tf.nn.conv2d(w_padded, w, self.stride, padding="VALID")
+
+    def _compute_target(self, w, convKxK_shape):
+        C_out = w.shape[-1]
+        outm3 = convKxK_shape[-3]
+        outm2 = convKxK_shape[-2]
+        ct = tf.cast(tf.math.floor(outm2 / 2), dtype=tf.int32)
+
+        target_zeros = tf.zeros((outm3 * outm2 - 1, C_out, C_out))
+        target = tf.concat(
+            [
+                target_zeros[: ct * outm2 + ct],
+                tf.expand_dims(tf.eye(C_out), axis=0),
+                target_zeros[ct * outm2 + ct :],
+            ],
+            axis=0,
+        )
+
+        target = tf.reshape(target, (outm3, outm2, C_out, C_out))
+        target = tf.transpose(target, [2, 0, 1, 3])
+        return target
+
+
+@register_keras_serializable("deel-lip", "LorthRegularizer")
+class LorthRegularizer(Regularizer):
+    def __init__(
+        self,
+        kernel_shape=None,
+        stride=1,
+        lambda_lorth=1.0,
+        dim=2,  # 2 for 2D conv, 1 for 1D conv
+        conv_transpose=False,
+    ) -> None:
+        """
+        Regularize a conv kernel to be orthogonal (all singular values are equal to 1)
+        using Lorth regularizer.
+
+        Args:
+            kernel_shape: the shape of the kernel.
+            stride (int): stride used in the associated convolution
+            lambda_lorth (float): weight of the orthogonalization regularization.
+            dim (int): 1 for 1D convolutions, 2 for 2D convolutions. Defaults to 2.
+            conv_transpose (bool): whether the kernel is from a transposed convolution.
+        """
+        super(LorthRegularizer, self).__init__()
+        self.kernel_shape = kernel_shape
+        self.stride = stride
+        self.lambda_lorth = lambda_lorth
+        self.dim = dim
+        self.conv_transpose = conv_transpose
+        if self.dim == 2:
+            self.lorth = Lorth2D(kernel_shape, stride, conv_transpose)
+        else:
+            raise NotImplementedError("Only 2D convolutions are supported for Lorth.")
+
+    def set_kernel_shape(self, shape):
+        self.kernel_shape = shape
+        self.lorth.set_kernel_shape(shape)
+
+    def __call__(self, x):
+        return self.lambda_lorth * self.lorth.compute_lorth(x)
+
+    def get_config(self):
+        return {
+            "kernel_shape": self.kernel_shape,
+            "stride": self.stride,
+            "lambda_lorth": self.lambda_lorth,
+            "dim": self.dim,
+            "conv_transpose": self.conv_transpose,
+        }
+
+
+@register_keras_serializable("deel-lip", "OrthDenseRegularizer")
+class OrthDenseRegularizer(Regularizer):
+    def __init__(self, lambda_orth=1.0) -> None:
+        """
+        Regularize a Dense kernel to be orthogonal (all singular values are equal to 1)
+        minimizing W.W^T-Id
+
+        Args:
+            lambda_orth (float): regularization factor (must be positive)
+        """
+        super(OrthDenseRegularizer, self).__init__()
+        self.lambda_orth = lambda_orth
+
+    def _dense_orth_dist(self, w):
+        transp_b = w.shape[0] <= w.shape[1]
+        # W.W^T if h<=w; W^T.W otherwise
+        wwt = tf.matmul(w, w, transpose_a=not transp_b, transpose_b=transp_b)
+        idx = tf.eye(wwt.shape[0])
+        return tf.reduce_sum(tf.square(wwt - idx))
+
+    def __call__(self, x):
+        return self.lambda_orth * self._dense_orth_dist(x)
+
+    def get_config(self):
+        return {"lambda_orth": self.lambda_orth}
diff --git a/doc/source/conf.py b/doc/source/conf.py
index eeef0cd4..46000190 100644
--- a/doc/source/conf.py
+++ b/doc/source/conf.py
@@ -14,7 +14,8 @@
 # import sys
 # sys.path.insert(0, os.path.abspath('.'))
 # import guzzle_sphinx_theme
-import sphinx_rtd_theme
+# import sphinx_rtd_theme
+from deel.lip import __version__
 
 # -- Project information -----------------------------------------------------
 
@@ -31,7 +32,8 @@
 )
 
 # The full version, including alpha/beta/rc tags
-release = "1.3.0"
+version = __version__
+release = version
 
 # -- General configuration ---------------------------------------------------
 
diff --git a/doc/source/deel.lip.regularizers.rst b/doc/source/deel.lip.regularizers.rst
new file mode 100644
index 00000000..ca676709
--- /dev/null
+++ b/doc/source/deel.lip.regularizers.rst
@@ -0,0 +1,8 @@
+deel.lip.regularizers module
+============================
+
+.. automodule:: deel.lip.regularizers
+   :exclude-members: Lorth, Lorth2D
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/doc/source/deel.lip.rst b/doc/source/deel.lip.rst
index f7664493..bcc405cc 100644
--- a/doc/source/deel.lip.rst
+++ b/doc/source/deel.lip.rst
@@ -16,6 +16,7 @@ Submodules
    deel.lip.metrics
    deel.lip.model
    deel.lip.normalizers
+   deel.lip.regularizers
    deel.lip.unconstrained_layers
    deel.lip.utils
 
diff --git a/setup.cfg b/setup.cfg
index 88b3d26b..87113135 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -4,18 +4,28 @@ max-line-length = 88
 extend-ignore =
     # See https://github.com/PyCQA/pycodestyle/issues/373
     E203, E231
-per-file-ignores = 
+per-file-ignores =
 	*/__init__.py: F401
 
 [tox:tox]
-envlist = py{36,37,38},py{36,37,38}-lint
+envlist = py{37,38,39,310}-tf{22,23,24,25,26,27,28,29,latest},py{37,38,39,310}-lint
 
 [testenv]
 deps =
+    packaging
+    tf22: protobuf == 3.20.3
+    tf22: tensorflow ~= 2.2.0
+    tf23: tensorflow ~= 2.3.0
+    tf24: tensorflow ~= 2.4.0
+    tf25: tensorflow ~= 2.5.0
+    tf26: tensorflow ~= 2.6.0
+    tf27: tensorflow ~= 2.7.0
+    tf28: tensorflow ~= 2.8.0
+    tf29: tensorflow ~= 2.9.0
 commands =
     python -m unittest
 
-[testenv:py{36,37,38}-lint]
+[testenv:py{37,38,39,310}-lint]
 skip_install = true
 deps =
     black
diff --git a/setup.py b/setup.py
index 88e729de..8ffbb484 100644
--- a/setup.py
+++ b/setup.py
@@ -11,6 +11,9 @@
 with open(path.join(this_directory, "README.md"), encoding="utf-8") as f:
     long_description = f.read()
 
+with open(path.join(this_directory, "deel/lip/VERSION")) as f:
+    version = f.read().strip()
+
 dev_requires = [
     "tox",
     "black",
@@ -30,7 +33,7 @@
 
 setuptools.setup(
     name="deel-lip",
-    version="1.3.0",
+    version=version,
     author=", ".join(["Mathieu SERRURIER", "Franck MAMALET", "Thibaut BOISSIN"]),
     author_email=", ".join(
         [
@@ -44,16 +47,18 @@
     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,<3"],
+    include_package_data=True,
+    install_requires=["numpy", "tensorflow~=2.2"],
     license="MIT",
     extras_require={"dev": dev_requires, "docs": docs_requires},
     classifiers=[
         "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
         "License :: OSI Approved :: MIT License",
         "Operating System :: OS Independent",
     ],
-    python_requires=">=3.6",
+    python_requires=">=3.7",
 )


Test patch:
diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml
index 3f6bbcf4..7f4f5fb6 100644
--- a/.github/workflows/python-tests.yml
+++ b/.github/workflows/python-tests.yml
@@ -1,6 +1,13 @@
 name: deel-lip tests
 
-on: [push, pull_request]
+on:
+  push:
+    branches:
+      - master
+      - develop
+  pull_request:
+  schedule:
+    - cron: "0 2 * * 0" # Run tests every Sunday at 2am
 
 jobs:
   checks:
@@ -8,17 +15,23 @@ jobs:
     strategy:
       max-parallel: 4
       matrix:
-        python-version: [3.6, 3.7, 3.8]
+        include:
+          - python-version: 3.7
+            tf-version: 2.3
+          - python-version: 3.9
+            tf-version: 2.7
+          - python-version: "3.10"
+            tf-version: latest
 
     steps:
-    - uses: actions/checkout@v1
+    - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v4
       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
+    - name: Test with tox (Python ${{ matrix.python-version }} - TensorFlow ${{ matrix.tf-version }})
+      run: tox -e py$(echo ${{ matrix.python-version }}-tf${{ matrix.tf-version }} | tr -d .)
diff --git a/tests/test_activations.py b/tests/test_activations.py
index 6411c007..c7bff996 100644
--- a/tests/test_activations.py
+++ b/tests/test_activations.py
@@ -4,7 +4,7 @@
 from tensorflow.keras.layers import Input
 from tensorflow.keras.optimizers import SGD
 from tensorflow.keras.losses import CategoricalCrossentropy
-from deel.lip.activations import GroupSort
+from deel.lip.activations import GroupSort, Householder
 import os
 import numpy as np
 
@@ -109,3 +109,184 @@ def test_idempotence(self):
         y1 = gs(x)
         y2 = gs(y1)
         np.testing.assert_equal(y1.numpy(), y2.numpy())
+
+
+class TestHouseholder(TestCase):
+    """Tests for Householder activation:
+    - instantiation of layer
+    - check outputs on dense (bs, n) tensor, with three thetas: 0, pi/2 and pi
+    - check outputs on dense (bs, h, w, n) tensor, with three thetas: 0, pi/2 and pi
+    - check idempotence hh(hh(x)) = hh(x)
+    """
+
+    def test_instantiation(self):
+        # Instantiation without argument
+        hh = Householder()
+        hh.build((28, 28, 10))
+        assert hh.theta.shape == (5,)
+        np.testing.assert_equal(hh.theta.numpy(), np.pi / 2)
+
+        # Instantiation with arguments
+        hh = Householder(
+            data_format="channels_last", k_coef_lip=2.5, theta_initializer="ones"
+        )
+        hh.build((32, 32, 16))
+        assert hh.theta.shape == (8,)
+        np.testing.assert_equal(hh.theta.numpy(), 1)
+
+        # Check serialization
+        hh = Householder(theta_initializer="glorot_uniform")
+        check_serialization(hh)
+
+        # Instantiation error because of wrong data format
+        with self.assertRaisesRegex(RuntimeError, "data format is supported"):
+            Householder(data_format="channels_first")
+
+    def test_theta_zero_dense_potentials(self):
+        """Householder with theta=0 on 2-D tensor (bs, n).
+        Theta=0 means Id if z2 > 0, and reflection if z2 < 0.
+        """
+        hh = Householder(theta_initializer="zeros")
+
+        bs = np.random.randint(64, 512)
+        n = np.random.randint(1, 1024) * 2
+
+        # Case 1: hh(x) = x   (identity case, z2 > 0)
+        z1 = tf.random.normal((bs, n // 2))
+        z2 = tf.random.uniform((bs, n // 2))
+        x = tf.concat([z1, z2], axis=-1)
+        np.testing.assert_allclose(hh(x), x)
+
+        # Case 2: hh(x) = [z1, -z2]   (reflection across z1 axis, z2 < 0)
+        z1 = tf.random.normal((bs, n // 2))
+        z2 = -tf.random.uniform((bs, n // 2))
+        x = tf.concat([z1, z2], axis=-1)
+        expected_output = tf.concat([z1, -z2], axis=-1)
+        np.testing.assert_allclose(hh(x), expected_output)
+
+    def test_theta_pi_dense_potentials(self):
+        """Householder with theta=pi on 2-D tensor (bs, n).
+        Theta=pi means Id if z1 < 0, and reflection if z1 > 0.
+        """
+        hh = Householder(theta_initializer=tf.keras.initializers.Constant(np.pi))
+
+        bs = np.random.randint(64, 512)
+        n = np.random.randint(1, 1024) * 2
+
+        # Case 1: hh(x) = x   (identity case, z1 < 0)
+        z1 = -tf.random.uniform((bs, n // 2))
+        z2 = tf.random.normal((bs, n // 2))
+        x = tf.concat([z1, z2], axis=-1)
+        np.testing.assert_allclose(hh(x), x, atol=1e-6)
+
+        # Case 2: hh(x) = [z1, -z2]   (reflection across z2 axis, z1 > 0)
+        z1 = tf.random.uniform((bs, n // 2))
+        z2 = tf.random.normal((bs, n // 2))
+        x = tf.concat([z1, z2], axis=-1)
+        expected_output = tf.concat([-z1, z2], axis=-1)
+        np.testing.assert_allclose(hh(x), expected_output, atol=1e-6)
+
+    def test_theta_90_dense_potentials(self):
+        """Householder with theta=pi/2 on 2-D tensor (bs, n).
+        Theta=pi/2 is equivalent to GroupSort2: Id if z1 < z2, and reflection if z1 > z2
+        """
+        hh = Householder()
+
+        bs = np.random.randint(64, 512)
+        n = np.random.randint(1, 1024) * 2
+
+        # Case 1: hh(x) = x   (identity case, z1 < z2)
+        z1 = -tf.random.normal((bs, n // 2))
+        z2 = z1 + tf.random.uniform((bs, n // 2))
+        x = tf.concat([z1, z2], axis=-1)
+        np.testing.assert_allclose(hh(x), x)
+
+        # Case 2: hh(x) = reflection(x)   (if z1 > z2)
+        z1 = tf.random.normal((bs, n // 2))
+        z2 = z1 - tf.random.uniform((bs, n // 2))
+        x = tf.concat([z1, z2], axis=-1)
+        expected_output = tf.concat([z2, z1], axis=-1)
+        np.testing.assert_allclose(hh(x), expected_output, atol=1e-6)
+
+    def test_theta_zero_conv_potentials(self):
+        """Householder with theta=0 on 4-D tensor (bs, h, w, c).
+        Theta=0 means Id if z2 > 0, and reflection if z2 < 0.
+        """
+        hh = Householder(theta_initializer="zeros")
+
+        bs = np.random.randint(32, 128)
+        h, w = np.random.randint(1, 64), np.random.randint(1, 64)
+        c = np.random.randint(1, 64) * 2
+
+        # Case 1: hh(x) = x   (identity case, z2 > 0)
+        z1 = tf.random.normal((bs, h, w, c // 2))
+        z2 = tf.random.uniform((bs, h, w, c // 2))
+        x = tf.concat([z1, z2], axis=-1)
+        np.testing.assert_allclose(hh(x), x)
+
+        # Case 2: hh(x) = [z1, -z2]   (reflection across z1 axis, z2 < 0)
+        z1 = tf.random.normal((bs, h, w, c // 2))
+        z2 = -tf.random.uniform((bs, h, w, c // 2))
+        x = tf.concat([z1, z2], axis=-1)
+        expected_output = tf.concat([z1, -z2], axis=-1)
+        np.testing.assert_allclose(hh(x), expected_output)
+
+    def test_theta_pi_conv_potentials(self):
+        """Householder with theta=pi on 4-D tensor (bs, h, w, c).
+        Theta=pi means Id if z1 < 0, and reflection if z1 > 0.
+        """
+        hh = Householder(theta_initializer=tf.keras.initializers.Constant(np.pi))
+
+        bs = np.random.randint(32, 128)
+        h, w = np.random.randint(1, 64), np.random.randint(1, 64)
+        c = np.random.randint(1, 64) * 2
+
+        # Case 1: hh(x) = x   (identity case, z1 < 0)
+        z1 = -tf.random.uniform((bs, h, w, c // 2))
+        z2 = tf.random.normal((bs, h, w, c // 2))
+        x = tf.concat([z1, z2], axis=-1)
+        np.testing.assert_allclose(hh(x), x, atol=1e-6)
+
+        # Case 2: hh(x) = [z1, -z2]   (reflection across z2 axis, z1 > 0)
+        z1 = tf.random.uniform((bs, h, w, c // 2))
+        z2 = tf.random.normal((bs, h, w, c // 2))
+        x = tf.concat([z1, z2], axis=-1)
+        expected_output = tf.concat([-z1, z2], axis=-1)
+        np.testing.assert_allclose(hh(x), expected_output, atol=1e-6)
+
+    def test_theta_90_conv_potentials(self):
+        """Householder with theta=pi/2 on 4-D tensor (bs, h, w, c).
+        Theta=pi/2 is equivalent to GroupSort2: Id if z1 < z2, and reflection if z1 > z2
+        """
+        hh = Householder()
+
+        bs = np.random.randint(32, 128)
+        h, w = np.random.randint(1, 64), np.random.randint(1, 64)
+        c = np.random.randint(1, 64) * 2
+
+        # Case 1: hh(x) = x   (identity case, z1 < z2)
+        z1 = -tf.random.normal((bs, h, w, c // 2))
+        z2 = z1 + tf.random.uniform((bs, h, w, c // 2))
+        x = tf.concat([z1, z2], axis=-1)
+        np.testing.assert_allclose(hh(x), x)
+
+        # Case 2: hh(x) = reflection(x)   (if z1 > z2)
+        z1 = tf.random.normal((bs, h, w, c // 2))
+        z2 = z1 - tf.random.uniform((bs, h, w, c // 2))
+        x = tf.concat([z1, z2], axis=-1)
+        expected_output = tf.concat([z2, z1], axis=-1)
+        np.testing.assert_allclose(hh(x), expected_output, atol=1e-6)
+
+    def test_idempotence(self):
+        """Assert idempotence of Householder activation: hh(hh(x)) = hh(x)"""
+        hh = Householder(theta_initializer="glorot_uniform")
+
+        bs = np.random.randint(32, 128)
+        h, w = np.random.randint(1, 64), np.random.randint(1, 64)
+        c = np.random.randint(1, 32) * 2
+        x = tf.random.normal((bs, h, w, c))
+
+        # Run two times the HH activation and compare both outputs
+        y = hh(x)
+        z = hh(y)
+        np.testing.assert_allclose(y, z)
diff --git a/tests/test_condense.py b/tests/test_condense.py
index f5b152e5..67e42ac5 100644
--- a/tests/test_condense.py
+++ b/tests/test_condense.py
@@ -11,6 +11,7 @@
 from deel.lip.layers import (
     SpectralDense,
     SpectralConv2D,
+    SpectralConv2DTranspose,
     FrobeniusDense,
     FrobeniusConv2D,
     ScaledL2NormPooling2D,
@@ -42,6 +43,7 @@ def test_as_supertype_sequential(self):
                 SpectralConv2D(2, (3, 3)),
                 ScaledL2NormPooling2D((2, 2)),
                 FrobeniusConv2D(2, (3, 3)),
+                SpectralConv2DTranspose(5, 3),
                 Flatten(),
                 Dense(4),
                 SpectralDense(4),
@@ -57,6 +59,7 @@ def test_as_supertype_model(self):
         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 = SpectralConv2DTranspose(5, 3, k_coef_lip=2.0)(x)
         x = Flatten()(x)
         x = Dense(4)(x)
         x = SpectralDense(4, k_coef_lip=2.0)(x)
diff --git a/tests/test_layers.py b/tests/test_layers.py
index 2a32d56b..5aca83a1 100644
--- a/tests/test_layers.py
+++ b/tests/test_layers.py
@@ -4,6 +4,7 @@
 # =====================================================================================
 import os
 import pprint
+import tempfile
 import unittest
 
 import numpy as np
@@ -30,6 +31,7 @@
     LipschitzLayer,
     SpectralDense,
     SpectralConv2D,
+    SpectralConv2DTranspose,
     FrobeniusDense,
     FrobeniusConv2D,
     ScaledAveragePooling2D,
@@ -187,7 +189,7 @@ def train_k_lip_model(
     np.random.seed(42)
     # create the keras model, defin opt, and compile it
     model = generate_k_lip_model(layer_type, layer_params, input_shape, k_lip_model)
-    optimizer = Adam(lr=0.001)
+    optimizer = Adam(learning_rate=0.001)
     model.compile(optimizer=optimizer, loss="mean_squared_error", metrics=[metrics.mse])
     # model.summary()
     # create the synthetic data generator
@@ -556,6 +558,49 @@ def test_spectralconv2d(self):
             ]
         )
 
+    def test_SpectralConv2DTranspose(self):
+        self._apply_tests_bank(
+            [
+                dict(
+                    layer_type=SpectralConv2DTranspose,
+                    layer_params={
+                        "filters": 2,
+                        "kernel_size": (3, 3),
+                        "use_bias": False,
+                    },
+                    batch_size=250,
+                    steps_per_epoch=125,
+                    epochs=5,
+                    input_shape=(5, 5, 1),
+                    k_lip_data=1.0,
+                    k_lip_model=1.0,
+                    callbacks=[],
+                ),
+                dict(
+                    layer_type=SpectralConv2DTranspose,
+                    layer_params={"filters": 2, "kernel_size": (3, 3)},
+                    batch_size=250,
+                    steps_per_epoch=125,
+                    epochs=5,
+                    input_shape=(5, 5, 1),
+                    k_lip_data=5.0,
+                    k_lip_model=1.0,
+                    callbacks=[],
+                ),
+                dict(
+                    layer_type=SpectralConv2DTranspose,
+                    layer_params={"filters": 2, "kernel_size": (3, 3)},
+                    batch_size=250,
+                    steps_per_epoch=125,
+                    epochs=5,
+                    input_shape=(5, 5, 1),
+                    k_lip_data=1.0,
+                    k_lip_model=5.0,
+                    callbacks=[],
+                ),
+            ]
+        )
+
     def test_frobeniusconv2d(self):
         # tests only checks that lip cons is enforced
         self._apply_tests_bank(
@@ -955,5 +1000,60 @@ def test_invertibleupsampling(self):
         )
 
 
+class TestSpectralConv2DTranspose(unittest.TestCase):
+    def test_instantiation(self):
+
+        # Supported cases
+        cases = (
+            dict(filters=5, kernel_size=3),
+            dict(filters=12, kernel_size=5, strides=2, use_bias=False),
+            dict(filters=3, kernel_size=3, padding="same", dilation_rate=1),
+            dict(filters=4, kernel_size=1, output_padding=None, activation="relu"),
+            dict(filters=16, kernel_size=3, data_format="channels_first"),
+        )
+
+        for i, kwargs in enumerate(cases):
+            with self.subTest(i=i):
+                SpectralConv2DTranspose(**kwargs)
+
+        # Unsupported cases
+        cases = (
+            {"msg": "Wrong padding", "kwarg": {"padding": "valid"}},
+            {"msg": "Wrong dilation rate", "kwarg": {"dilation_rate": 2}},
+            {"msg": "Wrong data format", "kwarg": {"output_padding": 5}},
+        )
+
+        for case in cases:
+            with self.subTest(case["msg"]):
+                with self.assertRaises(ValueError):
+                    SpectralConv2DTranspose(10, 3, **case["kwarg"])
+
+    def test_vanilla_export(self):
+        kwargs = dict(
+            filters=16,
+            kernel_size=5,
+            strides=2,
+            activation="relu",
+            data_format="channels_first",
+            input_shape=(28, 28, 3),
+        )
+
+        lay = SpectralConv2DTranspose(**kwargs)
+        model = Sequential([lay])
+
+        x = tf.random.normal((5,) + (kwargs["input_shape"]))
+        y1 = model(x)
+
+        # Test vanilla export inference comparison
+        vanilla_model = model.vanilla_export()
+        y2 = vanilla_model(x)
+        np.testing.assert_allclose(y1, y2, atol=1e-6)
+
+        # Test saving/loading model
+        with tempfile.TemporaryDirectory():
+            model.save("model.h5")
+            tf.keras.models.load_model("model.h5")
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_losses.py b/tests/test_losses.py
index a547dc97..53c23e1a 100644
--- a/tests/test_losses.py
+++ b/tests/test_losses.py
@@ -13,6 +13,8 @@
     MulticlassHinge,
     MulticlassHKR,
     MultiMargin,
+    TauCategoricalCrossentropy,
+    CategoricalHinge,
 )
 from deel.lip.utils import process_labels_for_multi_gpu
 import os
@@ -53,8 +55,8 @@ class Test(TestCase):
     def test_kr_loss(self):
         loss = KR()
 
-        y_true = tf.convert_to_tensor([1.0, 1.0, 1.0, 0.0, 0.0, 0.0])
-        y_pred = tf.convert_to_tensor([0.5, 1.5, -0.5, -0.5, -1.5, 0.5])
+        y_true = binary_tf_data([1.0, 1.0, 1.0, 0.0, 0.0, 0.0])
+        y_pred = binary_tf_data([0.5, 1.5, -0.5, -0.5, -1.5, 0.5])
         loss_val = loss(y_true, y_pred).numpy()
         self.assertEqual(loss_val, np.float32(1), "KR loss must be equal to 1")
 
@@ -79,9 +81,9 @@ def test_hkr_loss(self):
         check_serialization(1, loss)
 
     def test_hinge_margin_loss(self):
-        loss = HingeMargin(1.0)
-        y_true = tf.convert_to_tensor([1.0, 1.0, 1.0, 0.0, 0.0, 0.0])
-        y_pred = tf.convert_to_tensor([0.5, 1.5, -0.5, -0.5, -1.5, 0.5])
+        loss = HingeMargin(2.0)
+        y_true = binary_tf_data([1.0, 1.0, 1.0, 0.0, 0.0, 0.0])
+        y_pred = binary_tf_data([0.5, 1.5, -0.5, -0.5, -1.5, 0.5])
         loss_val = loss(y_true, y_pred).numpy()
         self.assertEqual(loss_val, np.float32(4 / 6), "Hinge loss must be equal to 4/6")
         loss_val_3 = loss(tf.cast(y_true, dtype=tf.int32), y_pred).numpy()
@@ -123,31 +125,8 @@ def test_kr_multiclass_loss(self):
         check_serialization(1, multiclass_kr)
 
     def test_hinge_multiclass_loss(self):
-        multiclass_hinge = MulticlassHinge(1.0)
-        hinge = HingeMargin(1.0)
-        y_true = tf.convert_to_tensor([1.0, 1.0, 1.0, 0.0, 0.0, 0.0])
-        y_pred = tf.convert_to_tensor([0.5, 1.5, -0.5, -0.5, -1.5, 0.5])
-        l_single = hinge(y_true, y_pred).numpy()
-        l_multi = multiclass_hinge(y_true, y_pred).numpy()
-        self.assertEqual(l_single, np.float32(4 / 6), "Hinge loss must be equal to 4/6")
-        self.assertEqual(
-            l_single,
-            l_multi,
-            "hinge multiclass must yield the same "
-            "results when given a single class "
-            "vector",
-        )
-        y_true = tf.expand_dims(y_true, -1)
-        y_pred = tf.expand_dims(y_pred, -1)
-        l_single = hinge(y_true, y_pred).numpy()
-        l_multi = multiclass_hinge(y_true, y_pred).numpy()
-        self.assertEqual(
-            l_single,
-            l_multi,
-            "hinge multiclass must yield the same "
-            "results when given a single class "
-            "vector",
-        )
+        multiclass_hinge = MulticlassHinge(2.0)
+
         n_class = 10
         n_items = 10000
         y_true = tf.one_hot(np.random.randint(0, 10, n_items), n_class)
@@ -165,7 +144,7 @@ def test_hinge_multiclass_loss(self):
         check_serialization(1, multiclass_hinge)
 
     def test_hkr_multiclass_loss(self):
-        multiclass_hkr = MulticlassHKR(5.0, 1.0)
+        multiclass_hkr = MulticlassHKR(5.0, 2.0)
         y_true = tf.one_hot([0, 0, 0, 1, 1, 2], 3)
         y_pred = np.float32(
             [
@@ -211,18 +190,46 @@ def test_multi_margin_loss(self):
         )
         check_serialization(1, multimargin_loss)
 
+    def test_categoricalhinge(self):
+        cathinge = CategoricalHinge(1.0)
+        n_class = 10
+        n_items = 10000
+        y_true = tf.one_hot(np.random.randint(0, 10, n_items), n_class)
+        y_pred = tf.random.normal((n_items, n_class))
+        loss_val = cathinge(y_true, y_pred).numpy()
+        loss_val_2 = cathinge(tf.cast(y_true, dtype=tf.int32), y_pred).numpy()
+        np.testing.assert_almost_equal(
+            loss_val_2, loss_val, 1, "test failed when y_true has dtype int32"
+        )
+        check_serialization(1, cathinge)
+
+    def test_tau_catcrossent(self):
+        taucatcrossent_loss = TauCategoricalCrossentropy(1.0)
+        n_class = 10
+        n_items = 10000
+        y_true = tf.one_hot(np.random.randint(0, 10, n_items), n_class)
+        y_pred = tf.random.normal((n_items, n_class))
+        loss_val = taucatcrossent_loss(y_true, y_pred).numpy()
+        loss_val_2 = taucatcrossent_loss(
+            tf.cast(y_true, dtype=tf.int32), y_pred
+        ).numpy()
+        np.testing.assert_almost_equal(
+            loss_val_2, loss_val, 1, "test failed when y_true has dtype int32"
+        )
+        check_serialization(1, taucatcrossent_loss)
+
     def test_no_reduction_binary_losses(self):
         """
         Assert binary losses without reduction. Three losses are tested on hardcoded
         y_true/y_pred of shape [8 elements, 1]: KR, HingeMargin and HKR.
         """
-        y_true = np.array([1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]).reshape((8, 1))
-        y_pred = np.array([0.5, 1.1, -0.1, 0.7, -1.3, -0.4, 0.2, -0.9]).reshape((8, 1))
+        y_true = binary_tf_data([1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0])
+        y_pred = binary_tf_data([0.5, 1.1, -0.1, 0.7, -1.3, -0.4, 0.2, -0.9])
 
         losses = (
             KR(reduction="none"),
-            HingeMargin(0.7, reduction="none"),
-            HKR(alpha=2.5, reduction="none"),
+            HingeMargin(0.7 * 2.0, reduction="none"),
+            HKR(alpha=2.5, min_margin=2.0, reduction="none"),
         )
 
         expected_loss_values = (
@@ -263,9 +270,11 @@ def test_no_reduction_multiclass_losses(self):
         )
         losses = (
             MulticlassKR(reduction="none"),
-            MulticlassHinge(reduction="none"),
+            MulticlassHinge(min_margin=2.0, reduction="none"),
             MultiMargin(0.7, reduction="none"),
-            MulticlassHKR(alpha=2.5, min_margin=0.5, reduction="none"),
+            MulticlassHKR(alpha=2.5, min_margin=1.0, reduction="none"),
+            CategoricalHinge(1.1, reduction="none"),
+            TauCategoricalCrossentropy(2.0, reduction="none"),
         )
 
         expected_loss_values = (
@@ -273,6 +282,19 @@ def test_no_reduction_multiclass_losses(self):
             np.float64([17, 13, 34, 15, 12, 62, 17, 45]) / 30,
             np.float64([0, 0, 19, 0, 0, 35, 0, 0]) / 30,
             np.float64([-779, -656, 1395, -625, -1609, 4557, -1284, 825]) / 900,
+            np.float64([0, 0.4, 2.3, 0.1, 0, 3.9, 0.4, 0]),
+            np.float64(
+                [
+                    0.044275,
+                    0.115109,
+                    1.243572,
+                    0.084923,
+                    0.010887,
+                    2.802300,
+                    0.114224,
+                    0.076357,
+                ]
+            ),
         )
 
         for loss, expected_loss_val in zip(losses, expected_loss_values):
@@ -301,8 +323,8 @@ def test_minibatches_binary_losses(self):
         reduction = "sum"
         losses = (
             KR(multi_gpu=True, reduction=reduction),
-            HingeMargin(0.7, reduction=reduction),
-            HKR(alpha=2.5, multi_gpu=True, reduction=reduction),
+            HingeMargin(0.7 * 2.0, reduction=reduction),
+            HKR(alpha=2.5, min_margin=2.0, multi_gpu=True, reduction=reduction),
         )
 
         expected_loss_values = (9.2, 2.2, 0.3)
@@ -389,10 +411,10 @@ def test_minibatches_multiclass_losses(self):
         reduction = "sum"
         losses = (
             MulticlassKR(multi_gpu=True, reduction=reduction),
-            MulticlassHinge(reduction=reduction),
+            MulticlassHinge(min_margin=2.0, reduction=reduction),
             MultiMargin(0.7, reduction=reduction),
             MulticlassHKR(
-                alpha=2.5, min_margin=0.5, multi_gpu=True, reduction=reduction
+                alpha=2.5, min_margin=1.0, multi_gpu=True, reduction=reduction
             ),
         )
 
@@ -434,7 +456,7 @@ def test_minibatches_multiclass_losses(self):
         num_items = 10000
         y_true = tf.one_hot(np.random.randint(num_classes, size=num_items), num_classes)
         y_true = process_labels_for_multi_gpu(y_true)
-        y_pred = tf.random.normal((num_items, num_classes))
+        y_pred = tf.random.normal((num_items, num_classes), seed=17)
 
         # Compare full batch loss and mini-batches loss
         for loss in losses:
diff --git a/tests/test_models.py b/tests/test_models.py
index 4dcecf18..674bce8d 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1,6 +1,7 @@
 """These tests assert that deel.lip Sequential and Model objects behave as expected."""
 
 import warnings
+from packaging import version
 from unittest import TestCase
 import numpy as np
 import tensorflow as tf
@@ -98,11 +99,13 @@ def test_warning_unsupported_1Lip_layers(self):
             kl.MaxPool2D(pool_size=3, strides=2),
             kl.Add(),
             kl.Concatenate(),
-            kl.Activation("gelu"),
             kl.Dense(5),
             kl.Conv2D(10, 3),
             kl.UpSampling2D(),
         ]
+        if version.parse(tf.__version__) >= version.parse("2.4.0"):
+            unsupported_layers.append(kl.Activation("gelu"))
+
         for lay in unsupported_layers:
             with self.assertWarnsRegex(
                 Warning,
diff --git a/tests/test_normalizers.py b/tests/test_normalizers.py
index 3cd273ad..ca1c798a 100644
--- a/tests/test_normalizers.py
+++ b/tests/test_normalizers.py
@@ -23,9 +23,8 @@ def test_spectral_normalization(self):
         kernel_shape = (15, 32)
         kernel = np.random.normal(size=kernel_shape).astype("float32")
         self._test_kernel(kernel)
-
-        # Conv kernel
-        kernel_shape = (5, 5, 64, 32)
+        # Dense kernel projection
+        kernel_shape = (32, 15)
         kernel = np.random.normal(size=kernel_shape).astype("float32")
         self._test_kernel(kernel)
 
diff --git a/tests/test_regularizers.py b/tests/test_regularizers.py
new file mode 100644
index 00000000..5ae30eb0
--- /dev/null
+++ b/tests/test_regularizers.py
@@ -0,0 +1,126 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+# =====================================================================================
+
+from unittest import TestCase
+
+import numpy as np
+import tensorflow as tf
+from deel.lip.regularizers import Lorth2D, LorthRegularizer
+
+
+class TestLorth2D(TestCase):
+    """Tests on Lorth2D class"""
+
+    def test_set_kernel_shape(self):
+        """Assert that set_kernel_shape() correctly sets parameters."""
+        # RO case
+        lorth = Lorth2D()
+        self.assertIsNone(lorth.kernel_shape)
+        self.assertIsNone(lorth.delta)
+        self.assertIsNone(lorth.padding)
+
+        kernel_shape = (5, 5, 32, 64)
+        lorth.set_kernel_shape(kernel_shape)
+        self.assertEqual(lorth.kernel_shape, kernel_shape)
+        self.assertEqual(lorth.delta, 32)
+        self.assertEqual(lorth.padding, 4)
+
+        # CO case
+        lorth = Lorth2D(stride=3)
+        kernel_shape = (5, 5, 8, 70)
+        lorth.set_kernel_shape(kernel_shape)
+
+        self.assertEqual(lorth.kernel_shape, kernel_shape)
+        self.assertEqual(lorth.delta, 0)  # zero, because CO case
+        self.assertEqual(lorth.padding, 3)
+
+    def test_existence_orthogonal_conv(self):
+        """Assert that Error and Warning are correctly raised when checking existence
+        of the orthogonal convolution."""
+
+        # RO case
+        with self.assertRaisesRegex(RuntimeError, "Impossible RO"):
+            kernel_shape = (3, 3, 32, 289)
+            Lorth2D(kernel_shape, stride=4)
+
+        # CO case
+        with self.assertRaisesRegex(RuntimeError, "Impossible CO"):
+            kernel_shape = (3, 3, 8, 129)
+            Lorth2D(kernel_shape, stride=4)
+
+        # square case (RO=CO)
+        with self.assertWarnsRegex(Warning, "hard to optimize"):
+            kernel_shape = (3, 3, 32, 32)
+            Lorth2D(kernel_shape, stride=1)
+
+    def test_compute_lorth(self):
+        """Assert Lorth2D computation on an identity kernel => must return 0"""
+
+        def _identity_kernel(kernel_size, channels):
+            assert kernel_size & 1
+            kernel_shape = (kernel_size, kernel_size, channels, channels)
+            w = np.zeros(kernel_shape, dtype=np.float32)
+            for i in range(channels):
+                w[kernel_size // 2, kernel_size // 2, i, i] = 1
+            return tf.constant(w)
+
+        kernel_size = (np.random.randint(1, 10) * 2) + 1
+        channels = np.random.randint(1, 100)
+        w = _identity_kernel(kernel_size, channels)
+
+        lorth = Lorth2D(w.shape)
+        loss = lorth.compute_lorth(w)
+        self.assertAlmostEqual(loss.numpy(), 0)
+
+
+class TestLorthRegularizer(TestCase):
+    """Tests on LorthRegularizer"""
+
+    def test_init(self):
+        """Check initialization from scratch and from config."""
+
+        kshape1 = (3, 3, 16, 32)
+        kshape2 = (5, 5, 16, 8)
+        stride = np.random.randint(1, 4)
+        lambda_lorth = np.random.normal()
+        regul = LorthRegularizer(kshape1, stride, lambda_lorth)
+
+        self.assertEqual(regul.lorth.kernel_shape, kshape1)
+
+        # Set new shape and get config
+        regul.set_kernel_shape(kshape2)
+        config = regul.get_config()
+        new_regul = LorthRegularizer.from_config(config)
+
+        self.assertEqual(new_regul.lorth.kernel_shape, kshape2)
+        self.assertEqual(new_regul.stride, stride)
+        self.assertEqual(new_regul.lambda_lorth, lambda_lorth)
+
+    def test_1D_not_supported(self):
+        """Assert that 1D convolutions are not supported."""
+        with self.assertRaisesRegex(NotImplementedError, "Only 2D convolutions"):
+            LorthRegularizer(dim=1)