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

Issue:
Vannilla export : Input layer is not duplicated
Within  the class Sequential vanilla_export does not export the Input layer 
Add Makefile for development
A Makefile is added to handle:

- dev environment creation
- testing (formatting / linting / unit tests)
- documentation generation

Moreover, the README is improved with logo and contributing section moved to CONTRIBUTING.md.
Element-wise losses for multi-GPU/TPU strategy
Keras losses API requires that subclasses return element-wise values, i.e. each element of a batch should return a single value. Then, a reduction operation performs either an average over the batch (default case) or a sum over the batch (useful for multi-GPU or TPU trainings).

This PR proposes a refactoring of all provided losses to fit the Keras losses API. For hinge losses, this is straightforward. For KR losses, it is more complicated: by definition, KR losses compute average over the batch. We then propose two element-wise KR formulations for standard training or multi-GPU/TPU training:
* a standard element-wise KR formulation that accepts classic one-hot encoded labels,
* a KR formulation appropriate for multi-GPU/TPU, where labels must be specifically pre-processed.

Multiple API changes are introduced:
* the KR loss is now subclassed from tf.keras Loss. Use `KR()(y_true, y_pred)` instead of `KR(y_true, y_pred)`
* the negative_KR function is removed. Use instead `HKR(alpha=0)`
* a new utils function `process_labels_for_KR_loss()` is added for label processing when using a multi-GPU/TPU strategy
* a new argument `multi_gpu` is added in all KR-based losses (KR, HKR, MulticlassKR, MulticlassHKR). When training with multi-GPU/TPU, this flag must be set to True and the labels must be processed using `process_labels_for_KR_loss()`.



Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 8664f7fd..418101ff 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,6 +7,7 @@ __pycache__
 # IDE / Tools files
 .vscode
 .tox
+deel_lip_dev_env
 
 # Files generated:
 logs
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
new file mode 100644
index 00000000..b52a1954
--- /dev/null
+++ b/CONTRIBUTING.md
@@ -0,0 +1,55 @@
+# Contributing
+
+Thanks for taking the time to contribute!
+
+From opening a bug report to creating a pull request: every contribution is
+appreciated and welcome. If you're planning to implement a new feature or change
+the API, please create an issue first. This way we can ensure that your precious
+work is not in vain.
+
+
+## Setup the development environment
+
+- Clone the repo and go to your local `deel-lip` folder
+- Create a virtual environment and install the development requirements:
+
+  `make prepare-dev && source deel_lip_dev_env/bin/activate`.
+- You are ready to install the library:
+
+  `pip install -e .`
+
+Welcome to the team!
+
+
+## Check your code changes
+
+Before opening a pull request, please make sure you check your code and you run the
+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.
+
+
+## Submitting your changes
+
+You can submit a pull request. Something that will increase the chance that your pull
+request is accepted:
+
+- Write tests and ensure that the existing ones pass.
+- If `make test` is successful, you have fair chances to pass the CI workflows (linting
+  and tests)
+- Write a [good commit message](https://tbaggery.com/2008/04/19/a-note-about-git-commit-messages.html) (we follow a lowercase convention).
+- For a major fix/feature make sure your PR has an issue and if it doesn't, please
+  create one. This would help discussion with the community, and polishing ideas in case
+  of a new feature.
diff --git a/Makefile b/Makefile
new file mode 100644
index 00000000..95fa8365
--- /dev/null
+++ b/Makefile
@@ -0,0 +1,34 @@
+.PHONY: help prepare-dev test test-disable-gpu doc ipynb-to-rst
+.DEFAULT: help
+
+help:
+	@echo "make prepare-dev"
+	@echo "       create and prepare development environment, use only once"
+	@echo "make test"
+	@echo "       run tests and linting on py36, py37, py38"
+	@echo "make test-disable-gpu"
+	@echo "       run test with gpu disabled"
+	@echo "make doc"
+	@echo "       build Sphinx docs documentation"
+	@echo "ipynb-to-rst"
+	@echo "       Transform notebooks to .rst files in documentation and generate the doc"
+
+prepare-dev:
+	python3 -m pip install virtualenv
+	python3 -m venv deel_lip_dev_env
+	. deel_lip_dev_env/bin/activate && pip install --upgrade pip
+	. deel_lip_dev_env/bin/activate && pip install -e .
+	. deel_lip_dev_env/bin/activate && pip install -e .[dev]
+	. deel_lip_dev_env/bin/activate && pip install -e .[docs]
+
+test:
+	. deel_lip_dev_env/bin/activate && tox
+
+test-disable-gpu:
+	. deel_lip_dev_env/bin/activate && CUDA_VISIBLE_DEVICES=-1 tox
+
+doc:
+	. deel_lip_dev_env/bin/activate && cd doc && make html && cd -
+
+ipynb-to-rst:
+	. deel_lip_dev_env/bin/activate && cd doc && ./generate_doc.sh && cd -
diff --git a/README.md b/README.md
index e616b6f7..e41d8b62 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,9 @@
-# deel-lip
+<img src="doc/source/logo.svg#gh-light-mode-only" alt="logo deel-lip" style="width:500px;"/>
+<img src="doc/source/logo_white.svg#gh-dark-mode-only" alt="logo deel-lip" style="width:500px;"/>
 
 [![Python](https://img.shields.io/pypi/pyversions/deel-lip.svg)](https://pypi.org/project/deel-lip)
 [![PyPI](https://img.shields.io/pypi/v/deel-lip.svg)](https://pypi.org/project/deel-lip)
+[![Downloads](https://pepy.tech/badge/deel-lip)](https://pepy.tech/project/deel-lip)
 [![Documentation](https://img.shields.io/badge/api-reference-blue.svg)](https://deel-lip.readthedocs.io)
 [![deel-lip tests](https://github.com/deel-ai/deel-lip/actions/workflows/python-tests.yml/badge.svg?branch=master)](https://github.com/deel-ai/deel-lip/actions/workflows/python-tests.yml)
 [![deel-lip linters](https://github.com/deel-ai/deel-lip/actions/workflows/python-linters.yml/badge.svg?branch=master)](https://github.com/deel-ai/deel-lip/actions/workflows/python-linters.yml)
@@ -146,8 +148,8 @@ In order to use `deel-lip`, you also need a [valid tensorflow installation](http
 ## Cite this work
 
 This library has been built to support the work presented in the paper
-[*Achieving robustness in classification using optimaltransport with Hinge regularization*](https://arxiv.org/abs/2006.06520)
-which aim provable and efficient robustness by design.
+[*Achieving robustness in classification using optimal transport with Hinge regularization*](https://arxiv.org/abs/2006.06520)
+which aims provable and efficient robustness by design.
 
 This work can be cited as:
 ```latex
@@ -161,20 +163,10 @@ Eprint = {arXiv:2006.06520},
 
 ## Contributing
 
-To contribute, you can open an [issue](https://github.com/deel-ai/deel-lip/issues), or fork this repository and then submit
-changes through a [pull-request](https://github.com/deel-ai/deel-lip/pulls).
-We use [`black`](https://pypi.org/project/black/) to format the code and follow PEP-8 convention. To check
-that your code will pass the lint-checks, you can run:
-
-```bash
-tox -e py36-lint
-```
-
-You need [`tox`](https://tox.readthedocs.io/en/latest/) in order to run this. You can install it via `pip`:
-
-```bash
-pip install tox
-```
+Contributions are welcome! You can open an [issue](https://github.com/deel-ai/deel-lip/issues),
+or fork this repository and then submit changes through a
+[pull-request](https://github.com/deel-ai/deel-lip/pulls).
+Take a look at our [contribution guidelines](CONTRIBUTING.md).
 
 ## License
 
diff --git a/deel/lip/__init__.py b/deel/lip/__init__.py
index 3b6f04a6..aea6c3bd 100644
--- a/deel/lip/__init__.py
+++ b/deel/lip/__init__.py
@@ -2,3 +2,14 @@
 # rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
 # CRIAQ and ANITI - https://www.deel.ai/
 # =====================================================================================
+from . import activations
+from . import callbacks
+from . import constraints
+from . import initializers
+from . import layers
+from . import unconstrained_layers
+from . import losses
+from . import metrics
+from .model import Sequential, Model, vanillaModel
+from . import normalizers
+from . import utils
diff --git a/deel/lip/activations.py b/deel/lip/activations.py
index 1e189bdb..e0efe12e 100644
--- a/deel/lip/activations.py
+++ b/deel/lip/activations.py
@@ -120,26 +120,33 @@ def __init__(
         self.data_format = data_format
 
     def build(self, input_shape):
+        input_shape = tf.TensorShape(input_shape)
         super(GroupSort, self).build(input_shape)
         self._init_lip_coef(input_shape)
         if (self.n is None) or (self.n > input_shape[self.channel_axis]):
             self.n = input_shape[self.channel_axis]
         if (input_shape[self.channel_axis] % self.n) != 0:
             raise RuntimeError("self.n has to be a divisor of the number of channels")
+        input_shape = tuple(input_shape.as_list())
+        self.flat_shape = (
+            (-1,) + input_shape[1:-1] + (input_shape[-1] // self.n, self.n)
+        )
+        self.out_shape = (-1,) + input_shape[1:]
 
     def _compute_lip_coef(self, input_shape=None):
         return 1.0
 
+    @tf.function
     def call(self, x, **kwargs):
-        fv = tf.reshape(x, [-1, self.n])
+        fv = tf.reshape(x, self.flat_shape)
         if self.n == 2:
-            b, c = tf.split(fv, 2, 1)
-            newv = tf.concat([tf.minimum(b, c), tf.maximum(b, c)], axis=1)
-            newv = tf.reshape(newv, tf.shape(x))
-            return newv
+            b, c = tf.split(fv, 2, -1)
+            newv = tf.concat([tf.minimum(b, c), tf.maximum(b, c)], axis=-1)
+            newv = tf.reshape(newv, self.out_shape)
+            return newv * self._get_coef()
 
         newv = tf.sort(fv)
-        newv = tf.reshape(newv, tf.shape(x))
+        newv = tf.reshape(newv, self.out_shape)
         return newv * self._get_coef()
 
     def get_config(self):
diff --git a/deel/lip/constraints.py b/deel/lip/constraints.py
index d438c3b9..cac8f735 100644
--- a/deel/lip/constraints.py
+++ b/deel/lip/constraints.py
@@ -9,7 +9,12 @@
 import tensorflow as tf
 from tensorflow.keras import backend as K
 from tensorflow.keras.constraints import Constraint
-from .normalizers import reshaped_kernel_orthogonalization
+from .normalizers import (
+    reshaped_kernel_orthogonalization,
+    DEFAULT_EPS_SPECTRAL,
+    DEFAULT_EPS_BJORCK,
+    DEFAULT_BETA_BJORCK,
+)
 from tensorflow.keras.utils import register_keras_serializable
 
 
@@ -71,7 +76,12 @@ def get_config(self):
 @register_keras_serializable("deel-lip", "SpectralConstraint")
 class SpectralConstraint(Constraint):
     def __init__(
-        self, k_coef_lip=1.0, niter_spectral=3, niter_bjorck=15, u=None
+        self,
+        k_coef_lip=1.0,
+        eps_spectral=DEFAULT_EPS_SPECTRAL,
+        eps_bjorck=DEFAULT_EPS_BJORCK,
+        beta_bjorck=DEFAULT_BETA_BJORCK,
+        u=None,
     ) -> None:
         """
         Ensure that *all* singular values of the weight matrix equals to 1. Computation
@@ -82,13 +92,15 @@ def __init__(
 
         Args:
             k_coef_lip: lipschitz coefficient of the weight matrix
-            niter_spectral: number of iteration to find the maximum singular value.
-            niter_bjorck: number of iteration with Bjorck algorithm..
+            eps_spectral: stopping criterion for the iterative power algorithm.
+            eps_bjorck: stopping criterion Bjorck algorithm.
+            beta_bjorck: beta parameter in bjorck algorithm.
             u: vector used for iterated power method, can be set to None (used for
                 serialization/deserialization purposes).
         """
-        self.niter_spectral = niter_spectral
-        self.niter_bjorck = niter_bjorck
+        self.eps_spectral = eps_spectral
+        self.eps_bjorck = eps_bjorck
+        self.beta_bjorck = beta_bjorck
         self.k_coef_lip = k_coef_lip
         if not (isinstance(u, tf.Tensor) or (u is None)):
             u = tf.convert_to_tensor(u)
@@ -100,16 +112,18 @@ def __call__(self, w):
             w,
             self.u,
             self.k_coef_lip,
-            self.niter_spectral,
-            self.niter_bjorck,
+            self.eps_spectral,
+            self.eps_bjorck,
+            self.beta_bjorck,
         )
         return wbar
 
     def get_config(self):
         config = {
             "k_coef_lip": self.k_coef_lip,
-            "niter_spectral": self.niter_spectral,
-            "niter_bjorck": self.niter_bjorck,
+            "eps_spectral": self.eps_spectral,
+            "eps_bjorck": self.eps_bjorck,
+            "beta_bjorck": self.beta_bjorck,
             "u": None if self.u is None else self.u.numpy(),
         }
         base_config = super(SpectralConstraint, self).get_config()
diff --git a/deel/lip/initializers.py b/deel/lip/initializers.py
index ae8083e6..18eb80cb 100644
--- a/deel/lip/initializers.py
+++ b/deel/lip/initializers.py
@@ -4,7 +4,12 @@
 # =====================================================================================
 from tensorflow.keras.initializers import Initializer, Orthogonal
 from tensorflow.keras import initializers
-from .normalizers import reshaped_kernel_orthogonalization
+from .normalizers import (
+    reshaped_kernel_orthogonalization,
+    DEFAULT_EPS_SPECTRAL,
+    DEFAULT_EPS_BJORCK,
+    DEFAULT_BETA_BJORCK,
+)
 from tensorflow.keras.utils import register_keras_serializable
 
 
@@ -12,8 +17,9 @@
 class SpectralInitializer(Initializer):
     def __init__(
         self,
-        niter_spectral=3,
-        niter_bjorck=15,
+        eps_spectral=DEFAULT_EPS_SPECTRAL,
+        eps_bjorck=DEFAULT_EPS_BJORCK,
+        beta_bjorck=DEFAULT_BETA_BJORCK,
         k_coef_lip=1.0,
         base_initializer=Orthogonal(gain=1.0, seed=None),
     ) -> None:
@@ -22,13 +28,16 @@ def __init__(
         normalization.
 
         Args:
-            niter_spectral: number of iteration to do with the iterative power method
-            niter_bjorck: number of iteration to do with the bjorck algorithm
+            eps_spectral: stopping criterion of iterative power method
+            eps_bjorck: float greater than 0, stopping criterion of
+                bjorck algorithm, setting it to None disable orthogonalization
+            beta_bjorck: beta parameter of bjorck algorithm
             base_initializer: method used to generate weights before applying the
                 orthonormalization
         """
-        self.niter_spectral = niter_spectral
-        self.niter_bjorck = niter_bjorck
+        self.eps_spectral = eps_spectral
+        self.eps_bjorck = eps_bjorck
+        self.beta_bjorck = beta_bjorck
         self.k_coef_lip = k_coef_lip
         self.base_initializer = initializers.get(base_initializer)
         super(SpectralInitializer, self).__init__()
@@ -39,15 +48,17 @@ def __call__(self, shape, dtype=None, partition_info=None):
             w,
             None,
             self.k_coef_lip,
-            self.niter_spectral,
-            self.niter_bjorck,
+            self.eps_spectral,
+            self.eps_bjorck,
+            self.beta_bjorck,
         )
         return wbar
 
     def get_config(self):
         return {
-            "niter_spectral": self.niter_spectral,
-            "niter_bjorck": self.niter_bjorck,
+            "eps_spectral": self.eps_spectral,
+            "eps_bjorck": self.eps_bjorck,
+            "beta_bjorck": self.beta_bjorck,
             "k_coef_lip": self.k_coef_lip,
             "base_initializer": initializers.serialize(self.base_initializer),
         }
diff --git a/deel/lip/layers.py b/deel/lip/layers.py
index 30668c3b..e95498ea 100644
--- a/deel/lip/layers.py
+++ b/deel/lip/layers.py
@@ -20,7 +20,7 @@
 By default the layers are 1 Lipschitz almost everywhere, which is efficient for
 wasserstein distance estimation. However for other problems (such as adversarial
 robustness) the user may want to use layers that are at most 1 lipschitz, this can
-be done by setting the param `niter_bjorck=0`.
+be done by setting the param `eps_bjorck=None`.
 """
 
 import abc
@@ -29,20 +29,13 @@
 import tensorflow as tf
 from tensorflow.keras import backend as K
 from tensorflow.keras.initializers import RandomNormal
-from tensorflow.keras.layers import (
-    Layer,
-    Dense,
-    Conv2D,
-    AveragePooling2D,
-    GlobalAveragePooling2D,
-)
+import tensorflow.keras.layers as keraslayers
 
 from .constraints import SpectralConstraint
 from .initializers import SpectralInitializer
 from .normalizers import (
-    DEFAULT_NITER_BJORCK,
-    DEFAULT_NITER_SPECTRAL,
-    DEFAULT_NITER_SPECTRAL_INIT,
+    DEFAULT_EPS_BJORCK,
+    DEFAULT_EPS_SPECTRAL,
     reshaped_kernel_orthogonalization,
     DEFAULT_BETA_BJORCK,
 )
@@ -158,16 +151,13 @@ def vanilla_export(self):
 
 
 @register_keras_serializable("deel-lip", "SpectralDense")
-class SpectralDense(Dense, LipschitzLayer, Condensable):
+class SpectralDense(keraslayers.Dense, LipschitzLayer, Condensable):
     def __init__(
         self,
         units,
         activation=None,
         use_bias=True,
-        kernel_initializer=SpectralInitializer(
-            niter_spectral=DEFAULT_NITER_SPECTRAL_INIT,
-            niter_bjorck=DEFAULT_NITER_BJORCK,
-        ),
+        kernel_initializer=SpectralInitializer(),
         bias_initializer="zeros",
         kernel_regularizer=None,
         bias_regularizer=None,
@@ -175,8 +165,8 @@ def __init__(
         kernel_constraint=None,
         bias_constraint=None,
         k_coef_lip=1.0,
-        niter_spectral=DEFAULT_NITER_SPECTRAL,
-        niter_bjorck=DEFAULT_NITER_BJORCK,
+        eps_spectral=DEFAULT_EPS_SPECTRAL,
+        eps_bjorck=DEFAULT_EPS_BJORCK,
         beta_bjorck=DEFAULT_BETA_BJORCK,
         **kwargs
     ):
@@ -205,8 +195,8 @@ def __init__(
                 the `kernel` weights matrix.
             bias_constraint: Constraint function applied to the bias vector.
             k_coef_lip: lipschitz constant to ensure
-            niter_spectral: number of iteration to find the maximum singular value.
-            niter_bjorck: number of iteration with Bjorck algorithm.
+            eps_spectral: stopping criterion for the iterative power algorithm.
+            eps_bjorck: stopping criterion Bjorck algorithm.
             beta_bjorck: beta parameter in bjorck algorithm.
 
         Input shape:
@@ -236,17 +226,21 @@ def __init__(
         )
         self._kwargs = kwargs
         self.set_klip_factor(k_coef_lip)
-        self.niter_spectral = niter_spectral
-        self.niter_bjorck = niter_bjorck
+        self.eps_spectral = eps_spectral
         self.beta_bjorck = beta_bjorck
-        if not ((self.beta_bjorck <= 0.5) and (self.beta_bjorck > 0.0)):
+        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.u = None
         self.sig = None
         self.wbar = None
         self.built = False
-        if self.niter_spectral < 1:
-            raise RuntimeError("niter_spectral has to be > 0")
+        if self.eps_spectral < 0:
+            raise RuntimeError("eps_spectral has to be > 0")
 
     def build(self, input_shape):
         super(SpectralDense, self).build(input_shape)
@@ -279,8 +273,8 @@ def call(self, x, training=True):
                 self.kernel,
                 self.u,
                 self._get_coef(),
-                self.niter_spectral,
-                self.niter_bjorck,
+                self.eps_spectral,
+                self.eps_bjorck,
                 self.beta_bjorck,
             )
             self.wbar.assign(wbar)
@@ -298,8 +292,8 @@ def call(self, x, training=True):
     def get_config(self):
         config = {
             "k_coef_lip": self.k_coef_lip,
-            "niter_spectral": self.niter_spectral,
-            "niter_bjorck": self.niter_bjorck,
+            "eps_spectral": self.eps_spectral,
+            "eps_bjorck": self.eps_bjorck,
             "beta_bjorck": self.beta_bjorck,
         }
         base_config = super(SpectralDense, self).get_config()
@@ -310,8 +304,8 @@ def condense(self):
             self.kernel,
             self.u,
             self._get_coef(),
-            self.niter_spectral,
-            self.niter_bjorck,
+            self.eps_spectral,
+            self.eps_bjorck,
             self.beta_bjorck,
         )
         self.kernel.assign(wbar)
@@ -320,7 +314,7 @@ def condense(self):
 
     def vanilla_export(self):
         self._kwargs["name"] = self.name
-        layer = Dense(
+        layer = keraslayers.Dense(
             units=self.units,
             activation=self.activation,
             use_bias=self.use_bias,
@@ -336,7 +330,7 @@ def vanilla_export(self):
 
 
 @register_keras_serializable("deel-lip", "SpectralConv2D")
-class SpectralConv2D(Conv2D, LipschitzLayer, Condensable):
+class SpectralConv2D(keraslayers.Conv2D, LipschitzLayer, Condensable):
     def __init__(
         self,
         filters,
@@ -347,10 +341,7 @@ def __init__(
         dilation_rate=(1, 1),
         activation=None,
         use_bias=True,
-        kernel_initializer=SpectralInitializer(
-            niter_spectral=DEFAULT_NITER_SPECTRAL_INIT,
-            niter_bjorck=DEFAULT_NITER_BJORCK,
-        ),
+        kernel_initializer=SpectralInitializer(),
         bias_initializer="zeros",
         kernel_regularizer=None,
         bias_regularizer=None,
@@ -358,8 +349,8 @@ def __init__(
         kernel_constraint=None,
         bias_constraint=None,
         k_coef_lip=1.0,
-        niter_spectral=DEFAULT_NITER_SPECTRAL,
-        niter_bjorck=DEFAULT_NITER_BJORCK,
+        eps_spectral=DEFAULT_EPS_SPECTRAL,
+        eps_bjorck=DEFAULT_EPS_BJORCK,
         beta_bjorck=DEFAULT_BETA_BJORCK,
         **kwargs
     ):
@@ -418,8 +409,8 @@ def __init__(
             kernel_constraint: Constraint function applied to the kernel matrix.
             bias_constraint: Constraint function applied to the bias vector.
             k_coef_lip: lipschitz constant to ensure
-            niter_spectral: number of iteration to find the maximum singular value.
-            niter_bjorck: number of iteration with Bjorck algorithm.
+            eps_spectral: stopping criterion for the iterative power algorithm.
+            eps_bjorck: stopping criterion Bjorck algorithm.
             beta_bjorck: beta parameter in bjorck algorithm.
 
         This documentation reuse the body of the original keras.layers.Conv2D doc.
@@ -455,13 +446,17 @@ def __init__(
         self.u = None
         self.sig = None
         self.wbar = None
-        self.niter_spectral = niter_spectral
+        self.eps_spectral = eps_spectral
         self.beta_bjorck = beta_bjorck
-        if not ((self.beta_bjorck <= 0.5) and (self.beta_bjorck > 0.0)):
+        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.niter_bjorck = niter_bjorck
-        if self.niter_spectral < 1:
-            raise RuntimeError("niter_spectral has to be > 0")
+        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")
 
     def build(self, input_shape):
         super(SpectralConv2D, self).build(input_shape)
@@ -510,21 +505,7 @@ def _compute_lip_coef(self, input_shape=None):
         else:
             sn1 = self.strides[0]
             sn2 = self.strides[1]
-            ho = np.floor(h / sn1)
-            wo = np.floor(w / sn2)
-            alphabar1 = np.floor(k1_div2 / sn1)
-            alphabar2 = np.floor(k2_div2 / sn2)
-            betabar1 = k1_div2 - alphabar1 * sn1
-            betabar2 = k2_div2 - alphabar2 * sn2
-            zl1 = (alphabar1 * sn1 + 2 * betabar1) * (alphabar1 + 1) / 2
-            zl2 = (alphabar2 * sn2 + 2 * betabar2) * (alphabar2 + 1) / 2
-            gamma1 = h - 1 - sn1 * np.ceil((h - 1 - k1_div2) / sn1)
-            gamma2 = w - 1 - sn2 * np.ceil((w - 1 - k2_div2) / sn2)
-            alphah1 = np.floor(gamma1 / sn1)
-            alphaw2 = np.floor(gamma2 / sn2)
-            zr1 = (alphah1 + 1) * (k1_div2 - gamma1 + sn1 * alphah1 / 2.0)
-            zr2 = (alphaw2 + 1) * (k2_div2 - gamma2 + sn2 * alphaw2 / 2.0)
-            coefLip = np.sqrt((h * w) / ((k1 * ho - zl1 - zr1) * (k2 * wo - zl2 - zr2)))
+            coefLip = np.sqrt(1.0 / (np.ceil(k1 / sn1) * np.ceil(k2 / sn2)))
         return coefLip
 
     def call(self, x, training=True):
@@ -533,8 +514,8 @@ def call(self, x, training=True):
                 self.kernel,
                 self.u,
                 self._get_coef(),
-                self.niter_spectral,
-                self.niter_bjorck,
+                self.eps_spectral,
+                self.eps_bjorck,
                 self.beta_bjorck,
             )
             self.wbar.assign(wbar)
@@ -559,8 +540,8 @@ def call(self, x, training=True):
     def get_config(self):
         config = {
             "k_coef_lip": self.k_coef_lip,
-            "niter_spectral": self.niter_spectral,
-            "niter_bjorck": self.niter_bjorck,
+            "eps_spectral": self.eps_spectral,
+            "eps_bjorck": self.eps_bjorck,
             "beta_bjorck": self.beta_bjorck,
         }
         base_config = super(SpectralConv2D, self).get_config()
@@ -571,8 +552,8 @@ def condense(self):
             self.kernel,
             self.u,
             self._get_coef(),
-            self.niter_spectral,
-            self.niter_bjorck,
+            self.eps_spectral,
+            self.eps_bjorck,
             self.beta_bjorck,
         )
         self.kernel.assign(wbar)
@@ -581,7 +562,7 @@ def condense(self):
 
     def vanilla_export(self):
         self._kwargs["name"] = self.name
-        layer = Conv2D(
+        layer = keraslayers.Conv2D(
             filters=self.filters,
             kernel_size=self.kernel_size,
             strides=self.strides,
@@ -602,7 +583,7 @@ def vanilla_export(self):
 
 
 @register_keras_serializable("deel-lip", "FrobeniusDense")
-class FrobeniusDense(Dense, LipschitzLayer, Condensable):
+class FrobeniusDense(keraslayers.Dense, LipschitzLayer, Condensable):
     """
     Identical and faster than a SpectralDense in the case of a single output. In the
     multi-neurons setting, this layer can be used:
@@ -619,10 +600,7 @@ def __init__(
         units,
         activation=None,
         use_bias=True,
-        kernel_initializer=SpectralInitializer(
-            niter_spectral=DEFAULT_NITER_SPECTRAL_INIT,
-            niter_bjorck=0,
-        ),
+        kernel_initializer=SpectralInitializer(),
         bias_initializer="zeros",
         kernel_regularizer=None,
         bias_regularizer=None,
@@ -696,7 +674,7 @@ def condense(self):
 
     def vanilla_export(self):
         self._kwargs["name"] = self.name
-        layer = Dense(
+        layer = keraslayers.Dense(
             units=self.units,
             activation=self.activation,
             use_bias=self.use_bias,
@@ -712,7 +690,7 @@ def vanilla_export(self):
 
 
 @register_keras_serializable("deel-lip", "FrobeniusConv2D")
-class FrobeniusConv2D(Conv2D, LipschitzLayer, Condensable):
+class FrobeniusConv2D(keraslayers.Conv2D, LipschitzLayer, Condensable):
     """
     Same as SpectralConv2D but in the case of a single output.
     """
@@ -727,10 +705,7 @@ def __init__(
         dilation_rate=(1, 1),
         activation=None,
         use_bias=True,
-        kernel_initializer=SpectralInitializer(
-            niter_spectral=DEFAULT_NITER_SPECTRAL_INIT,
-            niter_bjorck=0,
-        ),
+        kernel_initializer=SpectralInitializer(),
         bias_initializer="zeros",
         kernel_regularizer=None,
         bias_regularizer=None,
@@ -817,9 +792,7 @@ def get_config(self):
         return dict(list(base_config.items()) + list(config.items()))
 
     def condense(self):
-        wbar = (
-            self.kernel / tf.norm(self.kernel, axis=self.axis_norm) * self._get_coef()
-        )
+        wbar = self.kernel / tf.norm(self.kernel) * self._get_coef()
         self.kernel.assign(wbar)
 
     def vanilla_export(self):
@@ -829,7 +802,7 @@ def vanilla_export(self):
 
 
 @register_keras_serializable("deel-lip", "ScaledAveragePooling2D")
-class ScaledAveragePooling2D(AveragePooling2D, LipschitzLayer):
+class ScaledAveragePooling2D(keraslayers.AveragePooling2D, LipschitzLayer):
     def __init__(
         self,
         pool_size=(2, 2),
@@ -902,7 +875,7 @@ def _compute_lip_coef(self, input_shape=None):
         return np.sqrt(np.prod(np.asarray(self.pool_size)))
 
     def call(self, x, training=True):
-        return super(AveragePooling2D, self).call(x) * self._get_coef()
+        return super(keraslayers.AveragePooling2D, self).call(x) * self._get_coef()
 
     def get_config(self):
         config = {
@@ -913,7 +886,7 @@ def get_config(self):
 
 
 @register_keras_serializable("deel-lip", "ScaledL2NormPooling2D")
-class ScaledL2NormPooling2D(AveragePooling2D, LipschitzLayer):
+class ScaledL2NormPooling2D(keraslayers.AveragePooling2D, LipschitzLayer):
     def __init__(
         self,
         pool_size=(2, 2),
@@ -1022,7 +995,7 @@ def get_config(self):
 
 
 @register_keras_serializable("deel-lip", "ScaledGlobalL2NormPooling2D")
-class ScaledGlobalL2NormPooling2D(GlobalAveragePooling2D, LipschitzLayer):
+class ScaledGlobalL2NormPooling2D(keraslayers.GlobalAveragePooling2D, LipschitzLayer):
     def __init__(self, data_format=None, k_coef_lip=1.0, eps_grad_sqrt=1e-6, **kwargs):
         """
         Average pooling operation for spatial data, with a lipschitz bound. This
@@ -1109,7 +1082,7 @@ def get_config(self):
 
 
 @register_keras_serializable("deel-lip", "ScaledGlobalAveragePooling2D")
-class ScaledGlobalAveragePooling2D(GlobalAveragePooling2D, LipschitzLayer):
+class ScaledGlobalAveragePooling2D(keraslayers.GlobalAveragePooling2D, LipschitzLayer):
     def __init__(self, data_format=None, k_coef_lip=1.0, **kwargs):
         """Global average pooling operation for spatial data with Lipschitz bound.
 
@@ -1169,7 +1142,7 @@ def get_config(self):
 
 
 @register_keras_serializable("deel-lip", "InvertibleDownSampling")
-class InvertibleDownSampling(Layer):
+class InvertibleDownSampling(keraslayers.Layer):
     def __init__(
         self, pool_size, data_format="channels_last", name=None, dtype=None, **kwargs
     ):
@@ -1236,7 +1209,7 @@ def get_config(self):
 
 
 @register_keras_serializable("deel-lip", "InvertibleUpSampling")
-class InvertibleUpSampling(Layer):
+class InvertibleUpSampling(keraslayers.Layer):
     def __init__(
         self, pool_size, data_format="channels_last", name=None, dtype=None, **kwargs
     ):
diff --git a/deel/lip/losses.py b/deel/lip/losses.py
index e1427f13..836ac9f4 100644
--- a/deel/lip/losses.py
+++ b/deel/lip/losses.py
@@ -6,6 +6,7 @@
 This module contains losses used in wasserstein distance estimation. See
 https://arxiv.org/abs/2006.06520 for more information.
 """
+from functools import partial
 import numpy as np
 import tensorflow as tf
 from tensorflow.keras.losses import Loss
@@ -13,56 +14,109 @@
 from tensorflow.keras.utils import register_keras_serializable
 
 
-@tf.function
-@register_keras_serializable("deel-lip", "KR")
-def KR(y_true, y_pred):
-    r"""
-    Loss to estimate wasserstein-1 distance using Kantorovich-Rubinstein duality.
-    The Kantorovich-Rubinstein duality is formulated as following:
+@register_keras_serializable("deel-lip", "_kr")
+def _kr(y_true, y_pred, epsilon):
+    """Returns the element-wise KR loss.
+
+    `y_true` and `y_pred` must be of rank 2: (batch_size, 1) for binary classification
+    or (batch_size, C) for multilabel/multiclass classification (with C categories).
+    `y_true` labels should be either 1 and 0, or 1 and -1.
+    """
+    y_true = tf.cast(y_true, y_pred.dtype)
+    batch_size = tf.cast(tf.shape(y_true)[0], dtype=y_pred.dtype)
+    # Transform y_true into {1, 0} values
+    S1 = tf.cast(tf.equal(y_true, 1), y_pred.dtype)
+    num_elements_per_class = tf.reduce_sum(S1, axis=0)
 
-    .. math::
-        W_1(\mu, \nu) =
-        \sup_{f \in Lip_1(\Omega)} \underset{\textbf{x} \sim \mu}{\mathbb{E}}
-        \left[f(\textbf{x} )\right] -
-        \underset{\textbf{x}  \sim \nu}{\mathbb{E}} \left[f(\textbf{x} )\right]
+    pos = S1 / (num_elements_per_class + epsilon)
+    neg = (1 - S1) / (batch_size - num_elements_per_class + epsilon)
+    # Since element-wise KR terms are averaged by loss reduction later on, it is needed
+    # to multiply by batch_size here.
+    # 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(batch_size * y_pred * (pos - neg), axis=-1)
 
 
-    Where mu and nu stands for the two distributions, the distribution where the
-    label is 1 and the rest.
+@register_keras_serializable("deel-lip", "_kr_multi_gpu")
+def _kr_multi_gpu(y_true, y_pred):
+    """Returns the element-wise KR loss when computing with a multi-GPU/TPU strategy.
 
-    Returns:
-        Callable, the function to compute Wasserstein loss
+    `y_true` and `y_pred` can be either of shape (batch_size, 1) or
+    (batch_size, # classes).
 
+    When using this loss function, the labels `y_true` must be pre-processed with the
+    :func:`process_labels_for_multi_gpu()` function.
     """
     y_true = tf.cast(y_true, y_pred.dtype)
-    # create two boolean masks each selecting one distribution
-    S0 = tf.equal(y_true, 1)
-    S1 = tf.not_equal(y_true, 1)
-    # compute the KR dual representation
-    return tf.reduce_mean(tf.boolean_mask(y_pred, S0)) - tf.reduce_mean(
-        tf.boolean_mask(y_pred, S1)
-    )
+    # Since the information of batch size was included in `y_true` by
+    # `process_labels_for_multi_gpu()`, there is no need here to multiply by batch size.
+    # 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(y_pred * y_true, axis=-1)
 
 
-@tf.function
-@register_keras_serializable("deel-lip", "negative_KR")
-def negative_KR(y_true, y_pred):
-    r"""
-    Loss to compute the negative wasserstein-1 distance using Kantorovich-Rubinstein
-    duality. This allows the maximisation of the term using conventional optimizer.
+@register_keras_serializable("deel-lip", "KR")
+class KR(Loss):
+    def __init__(self, multi_gpu=False, reduction=Reduction.AUTO, name="KR"):
+        r"""
+        Loss to estimate Wasserstein-1 distance using Kantorovich-Rubinstein duality.
+        The Kantorovich-Rubinstein duality is formulated as following:
 
-    Returns:
-        Callable, the function to compute negative Wasserstein loss
+        .. math::
+            W_1(\mu, \nu) =
+            \sup_{f \in Lip_1(\Omega)} \underset{\textbf{x} \sim \mu}{\mathbb{E}}
+            \left[f(\textbf{x} )\right] -
+            \underset{\textbf{x}  \sim \nu}{\mathbb{E}} \left[f(\textbf{x} )\right]
 
-    """
-    return -KR(y_true, y_pred)
+
+        Where mu and nu stands for the two distributions, the distribution where the
+        label is 1 and the rest.
+
+        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.
+
+        Using a multi-GPU/TPU strategy requires to set `multi_gpu` to True and to
+        pre-process the labels `y_true` with the
+        :func:`deel.lip.utils.process_labels_for_multi_gpu()` function.
+
+        Args:
+            multi_gpu (bool): set to True when running on multi-GPU/TPU
+            reduction: passed to tf.keras.Loss constructor
+            name: passed to tf.keras.Loss constructor
+
+        """
+        self.eps = 1e-7
+        self.multi_gpu = multi_gpu
+        super(KR, self).__init__(reduction=reduction, name=name)
+        if multi_gpu:
+            self.kr_function = _kr_multi_gpu
+        else:
+            self.kr_function = partial(_kr, epsilon=self.eps)
+
+    @tf.function
+    def call(self, y_true, y_pred):
+        return self.kr_function(y_true, y_pred)
+
+    def get_config(self):
+        config = {"multi_gpu": self.multi_gpu}
+        base_config = super(KR, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
 
 
 @register_keras_serializable("deel-lip", "HKR")
 class HKR(Loss):
-    def __init__(self, alpha, min_margin=1.0, reduction=Reduction.AUTO, name="HKR"):
+    def __init__(
+        self,
+        alpha,
+        min_margin=1.0,
+        multi_gpu=False,
+        reduction=Reduction.AUTO,
+        name="HKR",
+    ):
         r"""
-        Wasserstein loss with a regularization param based on hinge loss.
+        Wasserstein loss with a regularization parameter based on the hinge margin loss.
 
         .. math::
             \inf_{f \in Lip_1(\Omega)} \underset{\textbf{x} \sim P_-}{\mathbb{E}}
@@ -71,38 +125,49 @@ def __init__(self, alpha, min_margin=1.0, reduction=Reduction.AUTO, name="HKR"):
             \underset{\textbf{x}}{\mathbb{E}} \left(\text{min_margin}
             -Yf(\textbf{x})\right)_+
 
+        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.
+
+        Using a multi-GPU/TPU strategy requires to set `multi_gpu` to True and to
+        pre-process the labels `y_true` with the
+        :func:`deel.lip.utils.process_labels_for_multi_gpu()` function.
+
         Args:
             alpha: regularization factor
             min_margin: minimal margin ( see hinge_margin_loss )
-            Kantorovich-rubinstein term of the loss. In order to be consistent
-            between hinge and KR, the first label must yield the positve class
-            while the second yields negative class.
+                Kantorovich-Rubinstein term of the loss. In order to be consistent
+                between hinge and KR, the first label must yield the positive class
+                while the second yields negative class.
+            multi_gpu (bool): set to True when running on multi-GPU/TPU
             reduction: passed to tf.keras.Loss constructor
             name: passed to tf.keras.Loss constructor
 
-        Returns:
-             a function that compute the regularized Wasserstein loss
-
         """
         self.alpha = alpha
         self.min_margin = min_margin
-        self.KR = KR
-        self.hinge = HingeMargin(min_margin)
+        self.multi_gpu = multi_gpu
+        self.KRloss = KR(multi_gpu=multi_gpu)
+        self.hingeloss = HingeMargin(min_margin)
         super(HKR, self).__init__(reduction=reduction, name=name)
 
     @tf.function
     def call(self, y_true, y_pred):
-        if self.alpha == np.inf:  # if alpha infinite, use hinge only
-            return self.hinge(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:
-            # true value: positive value should be the first to be coherent with the
-            # hinge loss (positive y_pred)
-            return self.alpha * self.hinge(y_true, y_pred) - self.KR(y_true, y_pred)
+            kr = -self.KRloss.call(y_true, y_pred)
+            hinge = self.hingeloss.call(y_true, y_pred)
+            return kr + self.alpha * hinge
 
     def get_config(self):
         config = {
             "alpha": self.alpha,
             "min_margin": self.min_margin,
+            "multi_gpu": self.multi_gpu,
         }
         base_config = super(HKR, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
@@ -118,24 +183,28 @@ def __init__(self, min_margin=1.0, reduction=Reduction.AUTO, name="HingeMargin")
             \underset{\textbf{x}}{\mathbb{E}} \left(\text{min_margin}
             -Yf(\textbf{x})\right)_+
 
+        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: the minimal margin to enforce.
+            min_margin: positive float, margin to enforce.
             reduction: passed to tf.keras.Loss constructor
             name: passed to tf.keras.Loss constructor
 
-        Returns:
-            a function that compute the hinge loss.
-
         """
         self.min_margin = min_margin
         super(HingeMargin, self).__init__(reduction=reduction, name=name)
 
     @tf.function
     def call(self, y_true, y_pred):
-        y_true = tf.where(y_true == 1, 1, -1)
-        sign = tf.cast(y_true, y_pred.dtype)
+        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)
-        return tf.reduce_mean(hinge)
+        # 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)
 
     def get_config(self):
         config = {
@@ -147,50 +216,41 @@ def get_config(self):
 
 @register_keras_serializable("deel-lip", "MulticlassKR")
 class MulticlassKR(Loss):
-    def __init__(self, reduction=Reduction.AUTO, name="MulticlassKR"):
+    def __init__(self, multi_gpu=False, reduction=Reduction.AUTO, name="MulticlassKR"):
         r"""
-        Loss to estimate average of W1 distance using Kantorovich-Rubinstein duality
-        over outputs. Note y_true should be one hot encoding (labels being 1s and 0s
-        ). In this multiclass setup thr KR term is computed for each class and then
-        averaged.
+        Loss to estimate average of Wasserstein-1 distance using Kantorovich-Rubinstein
+        duality over outputs. In this multiclass setup, the KR term is computed for each
+        class and then averaged.
+
+        Note that `y_true` should be one-hot encoded or pre-processed with the
+        :func:`deel.lip.utils.process_labels_for_multi_gpu()` function.
+
+        Using a multi-GPU/TPU strategy requires to set `multi_gpu` to True and to
+        pre-process the labels `y_true` with the
+        :func:`deel.lip.utils.process_labels_for_multi_gpu()` function.
 
         Args:
+            multi_gpu (bool): set to True when running on multi-GPU/TPU
             reduction: passed to tf.keras.Loss constructor
             name: passed to tf.keras.Loss constructor
 
-        Returns:
-            Callable, the function to compute Wasserstein multiclass loss.
-            #Note y_true has to be one hot encoded
-
         """
         self.eps = 1e-7
+        self.multi_gpu = multi_gpu
         super(MulticlassKR, self).__init__(reduction=reduction, name=name)
+        if multi_gpu:
+            self.kr_function = _kr_multi_gpu
+        else:
+            self.kr_function = partial(_kr, epsilon=self.eps)
 
     @tf.function
     def call(self, y_true, y_pred):
-        y_true = tf.cast(y_true, y_pred.dtype)
-        # use y_true to zero out y_pred where y_true != 1
-        # espYtrue is the avg value of y_pred when y_true==1
-        # (one average per output neuron)
-        espYtrue = tf.reduce_sum(y_pred * y_true, axis=0) / (
-            tf.reduce_sum(y_true, axis=0) + self.eps
-        )
-        # use(1- y_true) to zero out y_pred where y_true == 1
-        # espNotYtrue is the avg value of y_pred when y_true==0
-        # (one average per output neuron)
-        espNotYtrue = tf.reduce_sum(y_pred * (1 - y_true), axis=0) / (
-            (
-                tf.cast(tf.shape(y_true)[0], dtype="float32")
-                - tf.reduce_sum(y_true, axis=0)
-            )
-            + self.eps
-        )
-        # compute the differences to have the KR term for each output neuron,
-        # and compute the average over the classes
-        return tf.reduce_mean(-espNotYtrue + espYtrue)
+        return self.kr_function(y_true, y_pred)
 
     def get_config(self):
-        return super(MulticlassKR, self).get_config()
+        config = {"multi_gpu": self.multi_gpu}
+        base_config = super(MulticlassKR, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
 
 
 @register_keras_serializable("deel-lip", "MulticlassHinge")
@@ -199,30 +259,28 @@ def __init__(
         self, min_margin=1.0, reduction=Reduction.AUTO, name="MulticlassHinge"
     ):
         """
-        Loss to estimate the Hinge loss in a multiclass setup. It compute the
-        elementwise hinge term. Note that this formulation differs from the one
-        commonly found in tensorflow/pytorch (with marximise the difference between
+        Loss to estimate the Hinge loss in a multiclass setup. It computes the
+        element-wise Hinge term. Note that this formulation differs from the one
+        commonly found in tensorflow/pytorch (which maximises the difference between
         the two largest logits). This formulation is consistent with the binary
-        classification loss used in a multiclass fashion. Note y_true should be one
-        hot encoded. labels in (1,0)
+        classification loss used in a multiclass fashion.
+
+        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 to enforce.
             reduction: passed to tf.keras.Loss constructor
             name: passed to tf.keras.Loss constructor
 
-        Returns:
-            Callable, the function to compute multiclass Hinge loss
-            #Note y_true has to be one hot encoded
-
         """
         self.min_margin = min_margin
         super(MulticlassHinge, self).__init__(reduction=reduction, name=name)
 
     @tf.function
     def call(self, y_true, y_pred):
-        sign = tf.where(y_true == 1, 1.0, -1.0)
-        y_true = tf.cast(y_true, y_pred.dtype)
+        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
@@ -231,7 +289,7 @@ def call(self, y_true, y_pred):
         else:
             factor = 1.0
         hinge = tf.where(sign > 0, hinge * factor, hinge)
-        return tf.reduce_mean(hinge)
+        return tf.reduce_mean(hinge, axis=-1)
 
     def get_config(self):
         config = {
@@ -247,6 +305,7 @@ def __init__(
         self,
         alpha=10.0,
         min_margin=1.0,
+        multi_gpu=False,
         reduction=Reduction.AUTO,
         name="MulticlassHKR",
     ):
@@ -254,37 +313,44 @@ def __init__(
         The multiclass version of HKR. This is done by computing the HKR term over each
         class and averaging the results.
 
+        Note that `y_true` should be one-hot encoded or pre-processed with the
+        :func:`deel.lip.utils.process_labels_for_multi_gpu()` function.
+
+        Using a multi-GPU/TPU strategy requires to set `multi_gpu` to True and to
+        pre-process the labels `y_true` with the
+        :func:`deel.lip.utils.process_labels_for_multi_gpu()` function.
+
         Args:
             alpha: regularization factor
-            min_margin: minimal margin ( see Hinge_multiclass_loss )
+            min_margin: positive float, margin to enforce.
+            multi_gpu (bool): set to True when running on multi-GPU/TPU
             reduction: passed to tf.keras.Loss constructor
             name: passed to tf.keras.Loss constructor
 
-        Returns:
-            Callable, the function to compute HKR loss
-            #Note y_true has to be one hot encoded
         """
         self.alpha = alpha
         self.min_margin = min_margin
+        self.multi_gpu = multi_gpu
         self.hingeloss = MulticlassHinge(self.min_margin)
-        self.KRloss = MulticlassKR(reduction=reduction, name=name)
+        self.KRloss = MulticlassKR(multi_gpu=multi_gpu, name=name)
         super(MulticlassHKR, self).__init__(reduction=reduction, name=name)
 
     @tf.function
     def call(self, y_true, y_pred):
         if self.alpha == np.inf:  # alpha = inf => hinge only
-            return self.hingeloss(y_true, y_pred)
+            return self.hingeloss.call(y_true, y_pred)
         elif self.alpha == 0.0:  # alpha = 0 => KR only
-            return -self.KRloss(y_true, y_pred)
+            return -self.KRloss.call(y_true, y_pred)
         else:
-            a = -self.KRloss(y_true, y_pred)
-            b = self.hingeloss(y_true, y_pred)
-            return a + self.alpha * b
+            kr = -self.KRloss.call(y_true, y_pred)
+            hinge = self.hingeloss.call(y_true, y_pred)
+            return kr + self.alpha * hinge
 
     def get_config(self):
         config = {
             "alpha": self.alpha,
             "min_margin": self.min_margin,
+            "multi_gpu": self.multi_gpu,
         }
         base_config = super(MulticlassHKR, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
@@ -292,38 +358,34 @@ def get_config(self):
 
 @register_keras_serializable("deel-lip", "MultiMargin")
 class MultiMargin(Loss):
-    def __init__(self, min_margin=1, reduction=Reduction.AUTO, name="MultiMargin"):
+    def __init__(self, min_margin=1.0, reduction=Reduction.AUTO, name="MultiMargin"):
         """
-        Compute the mean hinge margin loss for multi class (equivalent to Pytorch
-         multi_margin_loss)
+        Compute the hinge margin loss for multiclass (equivalent to Pytorch
+        multi_margin_loss)
+
+        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: the minimal margin to enforce.
+            min_margin: positive float, margin to enforce.
             reduction: passed to tf.keras.Loss constructor
             name: passed to tf.keras.Loss constructor
 
-        Notes:
-            y_true has to be to_categorical
-
-        Returns:
-            Callable, the function to compute multi margin loss
         """
         self.min_margin = min_margin
         super(MultiMargin, self).__init__(reduction=reduction, name=name)
 
     @tf.function
     def call(self, y_true, y_pred):
-        y_true = tf.where(y_true == 1, 1, 0)
-        y_true = tf.cast(y_true, y_pred.dtype)
+        mask = tf.where(y_true > 0, 1, 0)
+        mask = tf.cast(mask, y_pred.dtype)
         # get the y_pred[target_class]
         # (zeroing out all elements of y_pred where y_true=0)
-        vYtrue = tf.reduce_sum(y_pred * y_true, axis=-1, keepdims=True)
+        vYtrue = tf.reduce_sum(y_pred * mask, axis=-1, keepdims=True)
         # computing elementwise margin term : margin + y_pred[i]-y_pred[target_class]
         margin = tf.nn.relu(self.min_margin - vYtrue + y_pred)
         # averaging on all outputs and batch
-        final_loss = tf.reduce_mean(
-            tf.where(y_true == 1, 0.0, margin)
-        )  # two steps is useless
+        final_loss = tf.reduce_mean((1.0 - mask) * margin, axis=-1)
         return final_loss
 
     def get_config(self):
diff --git a/deel/lip/metrics.py b/deel/lip/metrics.py
new file mode 100644
index 00000000..b2a6e441
--- /dev/null
+++ b/deel/lip/metrics.py
@@ -0,0 +1,292 @@
+# 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/
+# =====================================================================================
+"""
+This module contains metrics applicable in provable robustness. See
+https://arxiv.org/abs/2006.06520 and https://arxiv.org/abs/2108.04062 for more
+information.
+"""
+import math
+import tensorflow as tf
+from tensorflow.keras.losses import Loss
+from tensorflow.keras.losses import Reduction
+from tensorflow.keras.utils import register_keras_serializable
+
+
+def _delta_multiclass(y_true, y_pred):
+    r"""
+    Compute the non-normalized provable robustness factor in multiclass setup (need
+    to be adjusted with
+    lipschitz constant.
+
+    .. math::
+            \Delta(x) = f_l(x) - \max_{i \neq l} f_i(x)
+
+    Args:
+        y_true: true labels must be in {1,0} or in {1,-1} (no label smoothing allowed)
+        y_pred: network predictions
+
+    Returns: non-normalized provable robustness factor
+
+    """
+    ynl_shape = (-1, tf.shape(y_pred)[-1] - 1)
+    yl = tf.boolean_mask(y_pred, y_true > 0)
+    ynl = tf.reshape(
+        tf.boolean_mask(y_pred, y_true <= 0),
+        ynl_shape,
+    )
+    delta = yl - tf.reduce_max(ynl, axis=-1)
+    return delta
+
+
+def _delta_binary(y_true, y_pred):
+    r"""
+    Compute the non-normalized provable robustness factor in binary setup (need to be
+    adjusted with
+    lipschitz constant).
+
+    .. math::
+            \Delta(x) = f(x) \text{ if } l=1, -f(x) \text{ otherwise}
+
+    Args:
+        y_true: true labels must be in {1,0} or in {1,-1} (no label smoothing allowed)
+        y_pred: network predictions
+
+    Returns: non-normalized provable robustness factor
+
+    """
+    y_true = tf.sign(tf.cast(y_true, y_pred.dtype) - 1e-3)
+    return tf.multiply(y_true, y_pred)
+
+
+@register_keras_serializable("deel-lip", "CategoricalProvableRobustAccuracy")
+class CategoricalProvableRobustAccuracy(Loss):
+    def __init__(
+        self,
+        epsilon=36 / 255,
+        lip_const=1.0,
+        disjoint_neurons=True,
+        reduction=Reduction.AUTO,
+        name="CategoricalProvableRobustAccuracy",
+    ):
+        r"""
+
+        The accuracy that can be proved at a given epsilon.
+
+        Args:
+            epsilon: the metric will return the guaranteed accuracy for the radius
+                epsilon
+            lip_const: lipschitz constant of the network
+            disjoint_neurons: must be set to True if your model ends with a
+                FrobeniusDense layer with `disjoint_neurons` set to True. Set to False
+                otherwise
+            reduction: the recution method when training in a multi-gpu / TPU system
+            name: metrics name.
+        """
+        self.lip_const = lip_const
+        self.epsilon = epsilon
+        self.disjoint_neurons = disjoint_neurons
+        if disjoint_neurons:
+            self.certificate_factor = 2 * lip_const
+        else:
+            self.certificate_factor = math.sqrt(2) * lip_const
+        super(CategoricalProvableRobustAccuracy, self).__init__(reduction, name)
+
+    @tf.function
+    def call(self, y_true, y_pred):
+        return tf.cast(
+            (_delta_multiclass(y_true, y_pred) / self.certificate_factor)
+            > self.epsilon,
+            y_pred.dtype,
+        )
+
+    def get_config(self):
+        config = {
+            "epsilon": self.epsilon,
+            "lip_const": self.lip_const,
+            "disjoint_neurons": self.disjoint_neurons,
+        }
+        base_config = super(CategoricalProvableRobustAccuracy, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+@register_keras_serializable("deel-lip", "BinaryProvableRobustAccuracy")
+class BinaryProvableRobustAccuracy(Loss):
+    def __init__(
+        self,
+        epsilon=36 / 255,
+        lip_const=1.0,
+        reduction=Reduction.AUTO,
+        name="BinaryProvableRobustAccuracy",
+    ):
+        r"""
+
+        The accuracy that can be proved at a given epsilon.
+
+        Args:
+            epsilon: the metric will return the guaranteed accuracy for the radius
+                epsilon
+            lip_const: lipschitz constant of the network
+            reduction: the recution method when training in a multi-gpu / TPU system
+            name: metrics name.
+        """
+        self.lip_const = lip_const
+        self.epsilon = epsilon
+        super(BinaryProvableRobustAccuracy, self).__init__(reduction, name)
+
+    @tf.function
+    def call(self, y_true, y_pred):
+        return tf.cast(
+            (_delta_binary(y_true, y_pred) / self.lip_const) > self.epsilon,
+            y_pred.dtype,
+        )
+
+    def get_config(self):
+        config = {
+            "epsilon": self.epsilon,
+            "lip_const": self.lip_const,
+        }
+        base_config = super(BinaryProvableRobustAccuracy, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+@register_keras_serializable("deel-lip", "CategoricalProvableAvgRobustness")
+class CategoricalProvableAvgRobustness(Loss):
+    def __init__(
+        self,
+        lip_const=1.0,
+        disjoint_neurons=True,
+        negative_robustness=False,
+        reduction=Reduction.AUTO,
+        name="CategoricalProvableAvgRobustness",
+    ):
+        r"""
+
+        Compute the average provable robustness radius on the dataset.
+
+        .. math::
+            \mathbb{E}_{x \in D}\left[ \frac{\phi\left(\mathcal{M}_f(x)\right)}{
+            L_f}\right]
+
+        :math:`\mathcal{M}_f(x)` is a term that: is positive when x is correctly
+        classified and negative otherwise. In both case the value give the robustness
+        radius around x.
+
+        In the multiclass setup we have:
+
+        .. math::
+            \mathcal{M}_f(x) =f_l(x) - \max_{i \neq l} f_i(x)
+
+        Where :math:`D` is the dataset, :math:`l` is the correct label for x and
+        :math:`L_f` is the lipschitz constant of the network (:math:`L = 2 \times
+        \text{lip_const}` when `disjoint_neurons=True`, :math:`L = \sqrt{2} \times
+        \text{lip_const}` otherwise).
+
+        When `negative_robustness` is set to `True` misclassified elements count as
+        negative robustness (:math:`\phi` act as identity function), when set to
+        `False`,
+        misclassified elements yield a robustness radius of 0 ( :math:`\phi(x)=relu(
+        x)` ). The elements are not ignored when computing the mean in both cases.
+
+        This metric works for labels both in {1,0} and {1,-1}.
+
+        Args:
+            lip_const: lipschitz constant of the network
+            disjoint_neurons: must be set to True is your model ends with a
+                FrobeniusDense layer with `disjoint_neurons` set to True. Set to False
+                otherwise
+            reduction: the recution method when training in a multi-gpu / TPU system
+            name: metrics name.
+        """
+        self.lip_const = lip_const
+        self.disjoint_neurons = disjoint_neurons
+        self.negative_robustness = negative_robustness
+        if disjoint_neurons:
+            self.certificate_factor = 2 * lip_const
+        else:
+            self.certificate_factor = math.sqrt(2) * lip_const
+        if self.negative_robustness:
+            self.delta_correction = lambda delta: delta
+        else:
+            self.delta_correction = tf.nn.relu
+        super(CategoricalProvableAvgRobustness, self).__init__(reduction, name)
+
+    @tf.function
+    def call(self, y_true, y_pred):
+        return (
+            self.delta_correction(_delta_multiclass(y_true, y_pred))
+            / self.certificate_factor
+        )
+
+    def get_config(self):
+        config = {
+            "lip_const": self.lip_const,
+            "disjoint_neurons": self.disjoint_neurons,
+            "negative_robustness": self.negative_robustness,
+        }
+        base_config = super(CategoricalProvableAvgRobustness, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+@register_keras_serializable("deel-lip", "BinaryProvableAvgRobustness")
+class BinaryProvableAvgRobustness(Loss):
+    def __init__(
+        self,
+        lip_const=1.0,
+        negative_robustness=False,
+        reduction=Reduction.AUTO,
+        name="BinaryProvableAvgRobustness",
+    ):
+        r"""
+
+        Compute the average provable robustness radius on the dataset.
+
+        .. math::
+            \mathbb{E}_{x \in D}\left[ \frac{\phi\left(\mathcal{M}_f(x)\right)}{
+            L_f}\right]
+
+        :math:`\mathcal{M}_f(x)` is a term that: is positive when x is correctly
+        classified and negative otherwise. In both case the value give the robustness
+        radius around x.
+
+        In the binary classification setup we have:
+
+        .. math::
+            \mathcal{M}_f(x) = f(x) \text{ if } l=1, -f(x) \text{otherwise}
+
+        Where :math:`D` is the dataset, :math:`l` is the correct label for x and
+        :math:`L_f` is the lipschitz constant of the network..
+
+        When `negative_robustness` is set to `True` misclassified elements count as
+        negative robustness (:math:`\phi` act as identity function), when set to
+        `False`,
+        misclassified elements yield a robustness radius of 0 ( :math:`\phi(x)=relu(
+        x)` ). The elements are not ignored when computing the mean in both cases.
+
+        This metric works for labels both in {1,0} and {1,-1}.
+
+        Args:
+            lip_const: lipschitz constant of the network
+            reduction: the recution method when training in a multi-gpu / TPU system
+            name: metrics name.
+        """
+        self.lip_const = lip_const
+        self.negative_robustness = negative_robustness
+        if self.negative_robustness:
+            self.delta_correction = lambda delta: delta
+        else:
+            self.delta_correction = tf.nn.relu
+        super(BinaryProvableAvgRobustness, self).__init__(reduction, name)
+
+    @tf.function
+    def call(self, y_true, y_pred):
+        return self.delta_correction(_delta_binary(y_true, y_pred)) / self.lip_const
+
+    def get_config(self):
+        config = {
+            "lip_const": self.lip_const,
+            "negative_robustness": self.negative_robustness,
+        }
+        base_config = super(BinaryProvableAvgRobustness, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
diff --git a/deel/lip/model.py b/deel/lip/model.py
index b9c55d72..c72424b4 100644
--- a/deel/lip/model.py
+++ b/deel/lip/model.py
@@ -7,14 +7,34 @@
 for condensation and vanilla exportation.
 """
 import math
-from typing import Dict
 from warnings import warn
 import numpy as np
-from tensorflow import Tensor
 from tensorflow.keras import Sequential as KerasSequential, Model as KerasModel
-from tensorflow.keras.layers import Input, InputLayer
+from tensorflow.keras import activations as ka
+from tensorflow.keras import layers as kl
 from .layers import LipschitzLayer, Condensable
 from tensorflow.keras.utils import register_keras_serializable
+from tensorflow.keras.models import clone_model
+
+
+_msg_not_lip = "Sequential model contains a layer which is not a 1-Lipschitz layer: {}"
+
+
+def _is_supported_1lip_layer(layer):
+    """Return True if the Keras layer is 1-Lipschitz. Note that in some cases, the layer
+    is 1-Lipschitz for specific set of parameters.
+    """
+    supported_1lip_layers = (kl.Softmax, kl.Flatten, kl.Reshape)
+    if isinstance(layer, supported_1lip_layers):
+        return True
+    elif isinstance(layer, kl.MaxPool2D):
+        return True if layer.pool_size <= layer.strides else False
+    elif isinstance(layer, kl.ReLU):
+        return True if (layer.threshold == 0 and layer.negative_slope <= 1) else False
+    elif isinstance(layer, kl.Activation):
+        supported_activations = (ka.linear, ka.relu, ka.sigmoid, ka.tanh)
+        return True if layer.activation in supported_activations else False
+    return False
 
 
 @register_keras_serializable("deel-lip", "Sequential")
@@ -50,46 +70,30 @@ def set_klip_factor(self, klip_factor):
         for layer in self.layers:
             if isinstance(layer, LipschitzLayer):
                 layer.set_klip_factor(math.pow(klip_factor, 1 / nb_layers))
-            else:
-                warn(
-                    "Sequential model contains a layer wich is not a Lipschitz layer: {}".format(  # noqa: E501
-                        layer.name
-                    )
-                )
+            elif _is_supported_1lip_layer(layer) is not True:
+                warn(_msg_not_lip.format(layer.name))
 
     def _compute_lip_coef(self, input_shape=None):
         for layer in self.layers:
             if isinstance(layer, LipschitzLayer):
                 layer._compute_lip_coef(input_shape)
-            else:
-                warn(
-                    "Sequential model contains a layer wich is not a Lipschitz layer: {}".format(  # noqa: E501
-                        layer.name
-                    )
-                )
+            elif _is_supported_1lip_layer(layer) is not True:
+                warn(_msg_not_lip.format(layer.name))
 
     def _init_lip_coef(self, input_shape):
         for layer in self.layers:
             if isinstance(layer, LipschitzLayer):
                 layer._init_lip_coef(input_shape)
-            else:
-                warn(
-                    "Sequential model contains a layer wich is not a Lipschitz layer: {}".format(  # noqa: E501
-                        layer.name
-                    )
-                )
+            elif _is_supported_1lip_layer(layer) is not True:
+                warn(_msg_not_lip.format(layer.name))
 
     def _get_coef(self):
         global_coef = 1.0
         for layer in self.layers:
             if isinstance(layer, LipschitzLayer) and (global_coef is not None):
                 global_coef *= layer._get_coef()
-            else:
-                warn(
-                    "Sequential model contains a layer wich is not a Lipschitz layer: {}".format(  # noqa: E501
-                        layer.name
-                    )
-                )
+            elif _is_supported_1lip_layer(layer) is not True:
+                warn(_msg_not_lip.format(layer.name))
                 global_coef = None
         return global_coef
 
@@ -99,17 +103,7 @@ def condense(self):
                 layer.condense()
 
     def vanilla_export(self):
-        layers = list()
-        for layer in self.layers:
-            if isinstance(layer, Condensable):
-                layers.append(layer.vanilla_export())
-            else:
-                lay_cp = layer.__class__.from_config(layer.get_config())
-                lay_cp.build(layer.input.shape[1:])
-                lay_cp.set_weights(layer.get_weights())
-                layers.append(lay_cp)
-        model = KerasSequential(layers, self.name)
-        return model
+        return vanillaModel(self)
 
     def get_config(self):
         config = {"k_coef_lip": self.k_coef_lip}
@@ -141,58 +135,34 @@ def vanilla_export(self) -> KerasModel:
         Returns:
             A Keras model, identical to this model, but where condensable layers have
             been replaced with their vanilla equivalent (e.g. SpectralConv2D with
-             Conv2D).
+            Conv2D).
         """
-        # Dictionary that will map tensor names (from the current model) to tensors
-        # in the exported model.# We initialize the dictionary for inputs:
-        tensors: Dict[str, Tensor] = {}
-
-        # Initialize the dictionary with inputs:
-        tensors.update({inp.name: Input(shape=inp.shape[1:]) for inp in self.inputs})
-
-        for lay in self.layers:
-
-            # Skip input layers:
-            if isinstance(lay, InputLayer):
-                continue
-
-            # Condense+Export the layer if it is a non-vanilla layer, otherwise
-            # just copy the layer:
-            if isinstance(lay, Condensable):
-                lay_cp = lay.vanilla_export()
-            else:
-                # Duplicate layer (weights are not duplicated):
-                lay_cp = lay.__class__.from_config(lay.get_config())
-                lay_cp.build(lay.input_shape)
-                lay_cp.set_weights(lay.get_weights().copy())
-
-            # For each input nodes, we are going to create corresponding operations
-            # in the exported models:
-            for inode in range(len(lay.inbound_nodes)):
-                inputs = lay.get_input_at(inode)
-                outputs = lay.get_output_at(inode)
-
-                # Fetch the
-                if isinstance(inputs, list):
-                    inputs = [tensors[input.name] for input in inputs]
-                else:
-                    inputs = tensors[inputs.name]
-
-                # Retrieve outputs layers (for the exported layer):
-                moutputs = lay_cp(inputs)
-
-                # Add the output tensors to the dictionary, using the names from the
-                # original model:
-                if isinstance(outputs, list):
-                    for outi, mouti in zip(outputs, moutputs):
-                        tensors[outi.name] = mouti
-                else:
-                    tensors[outputs.name] = moutputs
-
-        return KerasModel(
-            [tensors[inp.name] for inp in self.inputs],
-            [tensors[out.name] for out in self.outputs],
-        )
-
-
-vanillaModel = Model.vanilla_export
+        return vanillaModel(self)
+
+
+def vanillaModel(model):
+    """
+    Transform a model to its equivalent "vanilla" model, i.e. a model where
+    `Condensable` layers are replaced with their vanilla equivalent. For example,
+    `SpectralConv2D` layers are converted to tf.keras `Conv2D` layers.
+
+    The input model can be a tf.keras Sequential/Model or a deel.lip Sequential/Model.
+
+    Args:
+        model: a tf.keras or deel.lip model with Condensable layers.
+
+    Returns:
+        A Keras model, identical to the input model where `Condensable` layers are
+        replaced with their vanilla counterparts.
+    """
+
+    def _replace_condensable_layer(layer):
+        # Return a vanilla layer if Condensable, else return a copy of the layer
+        if isinstance(layer, Condensable):
+            return layer.vanilla_export()
+        new_layer = layer.__class__.from_config(layer.get_config())
+        new_layer.build(layer.input_shape)
+        new_layer.set_weights(layer.get_weights())
+        return new_layer
+
+    return clone_model(model, clone_function=_replace_condensable_layer)
diff --git a/deel/lip/normalizers.py b/deel/lip/normalizers.py
index 209f1d8d..9aed3242 100644
--- a/deel/lip/normalizers.py
+++ b/deel/lip/normalizers.py
@@ -9,18 +9,17 @@
 import tensorflow as tf
 from tensorflow.keras import backend as K
 
-DEFAULT_NITER_BJORCK = 15
-DEFAULT_NITER_SPECTRAL = 3
-DEFAULT_NITER_SPECTRAL_INIT = 10
-DEFAULT_BETA_BJORCK = 0.25
+DEFAULT_BETA_BJORCK = 0.5
+DEFAULT_EPS_SPECTRAL = 1e-3
+DEFAULT_EPS_BJORCK = 1e-3
 
 
 def reshaped_kernel_orthogonalization(
     kernel,
     u,
     adjustment_coef,
-    niter_spectral=DEFAULT_NITER_SPECTRAL,
-    niter_bjorck=DEFAULT_NITER_BJORCK,
+    eps_spectral=DEFAULT_EPS_SPECTRAL,
+    eps_bjorck=DEFAULT_EPS_BJORCK,
     beta=DEFAULT_BETA_BJORCK,
 ):
     """
@@ -33,84 +32,124 @@ def reshaped_kernel_orthogonalization(
         kernel: the kernel to orthogonalize
         u: the vector used to do the power iteration method
         adjustment_coef: the adjustment coefficient as used in convolution
-        niter_spectral: number of iteration to do in spectral algorithm
-        niter_bjorck: iteration used for bjorck algorithm
+        eps_spectral: stopping criterion in spectral algorithm
+        eps_bjorck: stopping criterion in bjorck algorithm
         beta: the beta used in the 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, niter=niter_spectral)
-    W_bar = bjorck_normalization(W_bar, niter=niter_bjorck, beta=beta)
+    W_bar, u, sigma = spectral_normalization(kernel, u, eps=eps_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 = W_bar * adjustment_coef
     W_bar = K.reshape(W_bar, kernel.shape)
     return W_bar, u, sigma
 
 
-def bjorck_normalization(w, niter=DEFAULT_NITER_BJORCK, beta=DEFAULT_BETA_BJORCK):
+def bjorck_normalization(w, eps=DEFAULT_EPS_BJORCK, beta=DEFAULT_BETA_BJORCK):
     """
     apply Bjorck normalization on w.
 
     Args:
         w: weight to normalize, in order to work properly, we must have
             max_eigenval(w) ~= 1
-        niter: number of iterations
+        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]
 
     Returns:
         the orthonormal weights
 
     """
-    for i in range(niter):
+    # create a fake old_w that does'nt pass the loop condition
+    # it won't affect computation as the first action done in the loop overwrite it.
+    old_w = 10 * w
+    # define the loop condition
+
+    def cond(w, old_w):
+        return tf.linalg.norm(w - old_w) >= eps
+
+    # define the loop body
+    def body(w, old_w):
+        old_w = w
         w = (1 + beta) * w - beta * w @ tf.transpose(w) @ 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
+    )
     return w
 
 
-def _power_iteration(w, u, niter=DEFAULT_NITER_SPECTRAL):
+def _power_iteration(w, u, eps=DEFAULT_EPS_SPECTRAL):
     """
     Internal function that performs the power iteration algorithm.
 
     Args:
         w: weights matrix that we want to find eigen vector
         u: initialization of the eigen vector
-        niter: number of iteration, must be greater than 0
+        eps: epsilon stopping criterion: norm(ut - ut-1) must be less than eps
 
     Returns:
          u and v corresponding to the maximum eigenvalue
 
     """
+    # build _u and _v (_v is size of _u@tf.transpose(w), will be set on the first body
+    # iteration)
     _u = u
-    for i in range(niter):
+    _v = tf.zeros(u.shape[:-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.
+    _old_u = 10 * _u
+
+    # define the loop condition
+    def cond(_u, _v, old_u):
+        return tf.linalg.norm(_u - old_u) >= eps
+
+    # define the loop body
+    def body(_u, _v, _old_u):
+        _old_u = _u
         _v = tf.math.l2_normalize(_u @ tf.transpose(w))
         _u = tf.math.l2_normalize(_v @ w)
+        return _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
+    )
     return _u, _v
 
 
-def spectral_normalization(kernel, u, niter=DEFAULT_NITER_SPECTRAL):
+def spectral_normalization(kernel, u, eps=DEFAULT_EPS_SPECTRAL):
     """
     Normalize the kernel to have it's max eigenvalue == 1.
 
     Args:
         kernel: the kernel to normalize
         u: initialization for the max eigen vector
-        niter: number of iteration
+        eps: epsilon stopping criterion: norm(ut - ut-1) must be less than eps
 
     Returns:
         the normalized kernel w_bar, it's shape, the maximum eigen vector, and the
-        maximum eigen value
+        maximum singular value
 
     """
     W_shape = kernel.shape
     if u is None:
-        niter *= 2  # if u was not known increase number of iterations
-        u = tf.ones(shape=tuple([1, W_shape[-1]]))
+        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, niter)
-    # Calculate Sigma
+    _u, _v = _power_iteration(W_reshaped, u, eps)
+    # compute Sigma
     sigma = _v @ W_reshaped
     sigma = sigma @ tf.transpose(_u)
     # normalize it
-    W_bar = W_reshaped / sigma
+    # 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)
     return W_bar, _u, sigma
diff --git a/deel/lip/unconstrained_layers.py b/deel/lip/unconstrained_layers.py
new file mode 100644
index 00000000..1702d7a3
--- /dev/null
+++ b/deel/lip/unconstrained_layers.py
@@ -0,0 +1,141 @@
+import tensorflow as tf
+from tensorflow.keras.utils import register_keras_serializable
+
+from .layers import Condensable
+from .utils import _padding_circular
+
+
+@register_keras_serializable("deel-lip", "PadConv2D")
+class PadConv2D(tf.keras.layers.Conv2D, Condensable):
+    def __init__(
+        self,
+        filters,
+        kernel_size,
+        strides=(1, 1),
+        padding="same",
+        data_format=None,
+        dilation_rate=(1, 1),
+        activation=None,
+        use_bias=True,
+        kernel_initializer="glorot_uniform",
+        bias_initializer="zeros",
+        kernel_regularizer=None,
+        bias_regularizer=None,
+        activity_regularizer=None,
+        kernel_constraint=None,
+        bias_constraint=None,
+        **kwargs
+    ):
+        """
+        This class is a Conv2D Layer with parameterized padding.
+        Since Conv2D layer only supports `"same"` and `"valid"` padding, this layer will
+        enable other type of padding, such as `"constant"`, `"symmetric"`, `"reflect"`
+        or `"circular"`.
+
+        Warning:
+            The PadConv2D is not a Lipschitz layer and must not be directly used. This
+            must be used as a base class to create a Lipschitz layer with padding.
+
+        All arguments are the same as the original `Conv2D` except the `padding`
+        which is defined as following:
+
+        Args:
+            padding: one of `"same"`, `"valid"` `"constant"`, `"symmetric"`,
+                `"reflect"` or `"circular"` (case-insensitive).
+        """
+        self.pad = lambda x: x
+        self.old_padding = padding
+        self.internal_input_shape = None
+        if padding.lower() != "same":  # same is directly processed in Conv2D
+            padding = "valid"
+        super(PadConv2D, self).__init__(
+            filters=filters,
+            kernel_size=kernel_size,
+            strides=strides,
+            padding=padding,
+            data_format=data_format,
+            dilation_rate=dilation_rate,
+            activation=activation,
+            use_bias=use_bias,
+            kernel_initializer=kernel_initializer,
+            bias_initializer=bias_initializer,
+            kernel_regularizer=kernel_regularizer,
+            bias_regularizer=bias_regularizer,
+            activity_regularizer=activity_regularizer,
+            kernel_constraint=kernel_constraint,
+            bias_constraint=bias_constraint,
+            **kwargs
+        )
+        self._kwargs = kwargs
+        if self.old_padding.lower() in ["same", "valid"]:
+            self.pad = lambda x: x
+            self.padding_size = [0, 0]
+        if self.old_padding.lower() in ["constant", "reflect", "symmetric"]:
+            self.padding_size = [self.kernel_size[0] // 2, self.kernel_size[1] // 2]
+            paddings = [
+                [0, 0],
+                [self.padding_size[0], self.padding_size[0]],
+                [self.padding_size[1], self.padding_size[1]],
+                [0, 0],
+            ]
+            self.pad = lambda t: tf.pad(t, paddings, self.old_padding)
+        if self.old_padding.lower() == "circular":
+            self.padding_size = [self.kernel_size[0] // 2, self.kernel_size[1] // 2]
+            self.pad = lambda t: _padding_circular(t, self.padding_size)
+
+    def compute_padded_shape(self, input_shape, padding_size):
+        if isinstance(input_shape, tf.TensorShape):
+            internal_input_shape = input_shape.as_list()
+        else:
+            internal_input_shape = list(input_shape)
+
+        first_spatial_dim = 1 if self.data_format == "channels_last" else 2
+        for index, pad in enumerate(padding_size):
+            internal_input_shape[first_spatial_dim + index] += 2 * pad
+        return tf.TensorShape(internal_input_shape)
+
+    def build(self, input_shape):
+        self.internal_input_shape = self.compute_padded_shape(
+            input_shape, self.padding_size
+        )
+        super(PadConv2D, self).build(self.internal_input_shape)
+
+    def compute_output_shape(self, input_shape):
+        return super(PadConv2D, self).compute_output_shape(self.internal_input_shape)
+
+    def call(self, x):
+        x = self.pad(x)
+        return super(PadConv2D, self).call(x)
+
+    def get_config(self):
+        base_config = super(PadConv2D, self).get_config()
+        base_config["padding"] = self.old_padding
+        return base_config
+
+    def condense(self):
+        return
+
+    def vanilla_export(self):
+        self._kwargs["name"] = self.name
+        if self.old_padding.lower() in ["same", "valid"]:
+            layer_type = tf.keras.layers.Conv2D
+        else:
+            layer_type = PadConv2D
+        layer = layer_type(
+            filters=self.filters,
+            kernel_size=self.kernel_size,
+            strides=self.strides,
+            padding=self.old_padding,
+            data_format=self.data_format,
+            dilation_rate=self.dilation_rate,
+            activation=self.activation,
+            use_bias=self.use_bias,
+            kernel_initializer="glorot_uniform",
+            bias_initializer="zeros",
+            **self._kwargs
+        )
+        layer.build(self.input_shape)
+        layer.kernel.assign(self.kernel)
+        if self.use_bias:
+            layer.bias.assign(self.bias)
+        return layer
diff --git a/deel/lip/utils.py b/deel/lip/utils.py
index 11ab086b..8be3cab5 100644
--- a/deel/lip/utils.py
+++ b/deel/lip/utils.py
@@ -7,6 +7,7 @@
 """
 from typing import Generator, Tuple, Any
 import numpy as np
+import tensorflow as tf
 from tensorflow.keras import Model
 from tensorflow.keras import backend as K
 
@@ -70,3 +71,75 @@ def evaluate_lip_const(model: Model, x, eps=1e-4, seed=None):
     lip_cst = K.max(ndfx / ndx)
     print("lip cst: %.3f" % lip_cst)
     return lip_cst
+
+
+def _padding_circular(x, circular_paddings):
+    """Add circular padding to a 4-D tensor. Only channels_last is supported."""
+    if circular_paddings is None:
+        return x
+    w_pad, h_pad = circular_paddings
+    if w_pad > 0:
+        x = tf.concat((x[:, -w_pad:, :, :], x, x[:, :w_pad, :, :]), axis=1)
+    if h_pad > 0:
+        x = tf.concat((x[:, :, -h_pad:, :], x, x[:, :, :h_pad, :]), axis=2)
+    return x
+
+
+def _zero_upscale2D(x, strides):
+    stride_v = strides[0] * strides[1]
+    if stride_v == 1:
+        return x
+    output_shape = x.get_shape().as_list()[1:]
+    if strides[1] > 1:
+        output_shape[1] *= strides[1]
+        x = tf.expand_dims(x, 3)
+        fillz = tf.zeros_like(x)
+        fillz = tf.tile(fillz, [1, 1, 1, strides[1] - 1, 1])
+        x = tf.concat((x, fillz), axis=3)
+        x = tf.reshape(x, (-1,) + tuple(output_shape))
+    if strides[0] > 1:
+        output_shape[0] *= strides[0]
+        x = tf.expand_dims(x, 2)
+        fillz = tf.zeros_like(x)
+        fillz = tf.tile(fillz, [1, 1, strides[0] - 1, 1, 1])
+        x = tf.concat((x, fillz), axis=2)
+        x = tf.reshape(x, (-1,) + tuple(output_shape))
+    return x
+
+
+def _maybe_transpose_kernel(w, transpose=False):
+    """Transpose 4-D kernel: permutation of axes 2 and 3 + reverse axes 0 and 1."""
+    if not transpose:
+        return w
+    w_adj = tf.transpose(w, perm=[0, 1, 3, 2])
+    w_adj = w_adj[::-1, ::-1, :]
+    return w_adj
+
+
+@tf.function
+def process_labels_for_multi_gpu(labels):
+    """Process labels to be fed to any loss based on KR estimation with a multi-GPU/TPU
+    strategy.
+
+    When using a multi-GPU/TPU strategy, the flag `multi_gpu` in KR-based losses must be
+    set to True and the labels have to be pre-processed with this function.
+
+    For binary classification, the labels should be of shape [batch_size, 1].
+    For multiclass problems, the labels must be one-hot encoded (1 or 0) with shape
+    [batch_size, number of classes].
+
+    Args:
+        labels: tf.Tensor containing the labels
+    Returns:
+        labels processed for KR-based losses with multi-GPU/TPU strategy.
+    """
+    eps = 1e-7
+    labels = tf.cast(tf.where(labels > 0, 1, 0), labels.dtype)
+    batch_size = tf.cast(tf.shape(labels)[0], labels.dtype)
+    counts = tf.reduce_sum(labels, axis=0)
+
+    pos = labels / (counts + eps)
+    neg = (1 - labels) / (batch_size - counts + eps)
+    # Since element-wise KR terms are averaged by loss reduction later on, it is needed
+    # to multiply by batch_size here.
+    return batch_size * (pos - neg)
diff --git a/doc/notebooks/demo1.ipynb b/doc/notebooks/demo1.ipynb
index a8bf7ca8..e42759bf 100644
--- a/doc/notebooks/demo1.ipynb
+++ b/doc/notebooks/demo1.ipynb
@@ -2,6 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "## Demo 1: Wasserstein distance estimation on toy example\n",
     "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deel-ai/deel-lip/blob/master/doc/notebooks/demo1.ipynb)\n",
@@ -36,21 +37,31 @@
     "\n",
     "\n",
     "We will illustrate this on a synthetic image dataset where $W_1$ is known."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "# pip install deel-lip -qqq"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Matplotlib created a temporary config/cache directory at /tmp/matplotlib-_d1mx9in because the default path (/home/thibaut.boissin/.config/matplotlib) is not a writable directory; it is highly recommended to set the MPLCONFIGDIR environment variable to a writable directory, in particular to speed up the import of Matplotlib and to better support multiprocessing.\n",
+      "2021-09-08 18:20:20.330918: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n"
+     ]
+    }
+   ],
    "source": [
     "from datetime import datetime\n",
     "import os\n",
@@ -66,52 +77,43 @@
     "\n",
     "from deel.lip.layers import SpectralConv2D, SpectralDense, FrobeniusDense\n",
     "from deel.lip.activations import MaxMin, GroupSort, FullSort\n",
-    "from deel.lip.losses import KR, negative_KR\n",
+    "from deel.lip.losses import KR, HKR\n",
     "from deel.lip.model import Model"
-   ],
-   "outputs": [
-    {
-     "output_type": "stream",
-     "name": "stderr",
-     "text": [
-      "Matplotlib created a temporary config/cache directory at /tmp/matplotlib-_d1mx9in because the default path (/home/thibaut.boissin/.config/matplotlib) is not a writable directory; it is highly recommended to set the MPLCONFIGDIR environment variable to a writable directory, in particular to speed up the import of Matplotlib and to better support multiprocessing.\n",
-      "2021-09-08 18:20:20.330918: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n"
-     ]
-    }
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "### Parameters input images\n",
     "\n",
     "The synthetic dataset will be composed image with black or white squares allowing us to check if the computed wasserstein distance is correct. One distribution will be the set of black images, while the other will be the set of images with a square on it. these two distribution are diracs, and the wasserstein distance can be analyticaly computed:\n",
     "\n",
     "In the case to the two diracs the wasserstein distance is then the L1 distance between the two images."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "img_size = 64 \n",
     "frac_value = 0.3  # proportion of the center square"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "### Generate images"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "def generate_toy_images(shape,frac=0,v=1):\n",
     "    \"\"\"\n",
@@ -162,13 +164,13 @@
     "        #np.random.shuffle(indexes_shuffle)\n",
     "        #yield  batch_x[indexes_shuffle,], batch_y[indexes_shuffle,]\n",
     "        yield  batch_x, batch_y"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "def display_img(img):\n",
     "    \"\"\"\n",
@@ -180,42 +182,31 @@
     "    \n",
     "    imgplot = ax.imshow((img*255).astype(np.uint))\n",
     "    "
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "Now let's take a look at the generated batches\n",
     "\n",
     "#### for binary generator"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 5,
-   "source": [
-    "test=binary_generator(2,(img_size,img_size,1),frac=frac_value)\n",
-    "imgs, y=next(test)\n",
-    "\n",
-    "display_img(imgs[0])\n",
-    "display_img(imgs[1])\n",
-    "print(\"Norm L2 \"+str(np.linalg.norm(imgs[1])))\n",
-    "print(\"Norm L2(count pixels) \"+str(math.sqrt(np.size(imgs[1][imgs[1]==1]))))"
-   ],
+   "metadata": {},
    "outputs": [
     {
-     "output_type": "stream",
      "name": "stdout",
+     "output_type": "stream",
      "text": [
       "Norm L2 35.0\n",
       "Norm L2(count pixels) 35.0\n"
      ]
     },
     {
-     "output_type": "display_data",
      "data": {
       "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAMb0lEQVR4nO3dT4yc9X3H8fenNi5pQmMbUsuyoQaBgjgEE1kUFFQRV0RuGgUOCBGlklOh7iWViFopgVZqm0qVyiWEQ1XJAhof2gAlTWz5UOI4RO3JYP4lBsfBSUHYsnErYyXpAdXw7WGebRdr1zuemWfG5fd+SdbO8+zsPl8x+97nmdnheVJVSHr/+5VZDyBpOoxdaoSxS40wdqkRxi41wtilRowVe5JtSQ4nOZLkvkkNJWnyMurf2ZOsAH4C3AYcBZ4FPldVr0xuPEmTsnKMr70ROFJVPwNI8hhwO7Bk7El8B4/Us6rKYuvHOYzfALyxYPlot07SBWicPftQkswBc31vR9K5jRP7MeDyBcsbu3XvUVU7gB3gYbw0S+Mcxj8LXJPkyiSrgLuB3ZMZS9Kkjbxnr6ozSf4IeApYATxaVS9PbDJJEzXyn95G2piH8VLv+ng1XtL/I8YuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcvGnuTRJCeTHFywbm2SvUle7T6u6XdMSeMaZs/+DWDbWevuA/ZV1TXAvm5Z0gVs2dir6l+BU2etvh3Y2d3eCdwx2bEkTdqoz9nXVdXx7vYJYN2E5pHUk5Ev2TyvqupcV2dNMgfMjbsdSeMZdc/+ZpL1AN3Hk0vdsap2VNWWqtoy4rYkTcCose8Gtne3twO7JjOOpL6kaskj8MEdkm8CtwKXAW8CfwF8B3gCuAJ4Hbirqs5+EW+x73XujUkaW1VlsfXLxj5Jxi71b6nYfQed1Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71IhlY09yeZKnk7yS5OUk93br1ybZm+TV7uOa/seVNKphrvW2HlhfVc8nuQR4DrgD+AJwqqr+Jsl9wJqq+soy38vLP0k9G/nyT1V1vKqe727/AjgEbABuB3Z2d9vJ4BeApAvUeT1nT7IJuAHYD6yrquPdp04A6yY7mqRJWjnsHZN8CPgW8KWq+nnyf0cKVVVLHaInmQPmxh1U0niGumRzkouAPcBTVfW1bt1h4NaqOt49r/9BVX10me/jc3apZyM/Z89gF/4IcGg+9M5uYHt3ezuwa9whJfVnmFfjbwH+DfgR8G63+k8ZPG9/ArgCeB24q6pOLfO93LNLPVtqzz7UYfykGLvUv5EP4yW9Pxi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRgxzrbeLkzyT5KUkLyf5arf+yiT7kxxJ8niSVf2PK2lUw+zZ3wa2VtX1wGZgW5KbgAeAB6vqauAt4J7eppQ0tmVjr4FfdosXdf8K2Ao82a3fCdzRx4CSJmOo5+xJViR5ETgJ7AV+CpyuqjPdXY4CG3qZUNJEDBV7Vb1TVZuBjcCNwLXDbiDJXJIDSQ6MNqKkSTivV+Or6jTwNHAzsDrJyu5TG4FjS3zNjqraUlVbxhlU0niGeTX+I0lWd7c/ANwGHGIQ/Z3d3bYDu3qaUdIEpKrOfYfkYwxegFvB4JfDE1X1V0muAh4D1gIvAL9fVW8v873OvTFJY6uqLLZ+2dgnydil/i0Vu++gkxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71AhjlxoxdOzdZZtfSLKnW74yyf4kR5I8nmRVf2NKGtf57NnvZXBBx3kPAA9W1dXAW8A9kxxM0mQNFXuSjcDvAQ93ywG2Ak92d9kJ3NHDfJImZNg9+9eBLwPvdsuXAqer6ky3fBTYMNnRJE3SMNdn/wxwsqqeG2UDSeaSHEhyYJSvlzQZK4e4zyeAzyb5NHAx8OvAQ8DqJCu7vftG4NhiX1xVO4Ad4CWbpVlads9eVfdX1caq2gTcDXy/qj4PPA3c2d1tO7CrtykljW2cv7N/BfjjJEcYPId/ZDIjSepDqqZ3ZO1hvNS/qspi630HndQIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9SIYS7sSJLXgF8A7wBnqmpLkrXA48Am4DXgrqp6q58xJY3rfPbsn6yqzVW1pVu+D9hXVdcA+7plSReocQ7jbwd2drd3AneMPY2k3gwbewHfTfJckrlu3bqqOt7dPgGsm/h0kiZmqOfswC1VdSzJbwB7k/x44Serqpa6Qmv3y2Fusc9Jmp7zvmRzkr8Efgn8IXBrVR1Psh74QVV9dJmv9ZLNUs9GvmRzkg8muWT+NvAp4CCwG9je3W07sGsyo0rqw7J79iRXAd/uFlcC/1hVf53kUuAJ4ArgdQZ/eju1zPdyzy71bKk9+3kfxo/D2KX+jXwYL+n9wdilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41YqjYk6xO8mSSHyc5lOTmJGuT7E3yavdxTd/DShrdsHv2h4B/qaprgeuBQ8B9wL6qugbY1y1LukANc2HHDwMvAlfVgjsnOYyXbJYuOONc6+1K4D+Av0/yQpKHu0s3r6uq4919TgDrJjOqpD4ME/tK4OPA31XVDcB/cdYhe7fHX3SvnWQuyYEkB8YdVtLohon9KHC0qvZ3y08yiP/N7vCd7uPJxb64qnZU1Zaq2jKJgSWNZtnYq+oE8EaS+efjvwO8AuwGtnfrtgO7eplQ0kQs+wIdQJLNwMPAKuBnwB8w+EXxBHAF8DpwV1WdWub7+AKd1LOlXqAbKvZJMXapf+O8Gi/pfcDYpUYYu9QIY5caYexSI4xdaoSxS41YOeXt/SeDN+Bc1t2epQthBnCOsznHe53vHL+51Cem+qaa/91ocmDW75W/EGZwDueY5hwexkuNMHapEbOKfceMtrvQhTADOMfZnOO9JjbHTJ6zS5o+D+OlRkw19iTbkhxOciTJ1M5Gm+TRJCeTHFywbuqnwk5yeZKnk7yS5OUk985iliQXJ3kmyUvdHF/t1l+ZZH/3+DyeZFWfcyyYZ0V3fsM9s5ojyWtJfpTkxflTqM3oZ6S307ZPLfYkK4C/BX4XuA74XJLrprT5bwDbzlo3i1NhnwH+pKquA24Cvtj9N5j2LG8DW6vqemAzsC3JTcADwINVdTXwFnBPz3PMu5fB6cnnzWqOT1bV5gV/6prFz0h/p22vqqn8A24GnlqwfD9w/xS3vwk4uGD5MLC+u70eODytWRbMsAu4bZazAL8GPA/8FoM3b6xc7PHqcfsbux/grcAeIDOa4zXgsrPWTfVxAT4M/Dvda2mTnmOah/EbgDcWLB/t1s3KTE+FnWQTcAOwfxazdIfOLzI4Uehe4KfA6ao6091lWo/P14EvA+92y5fOaI4CvpvkuSRz3bppPy69nrbdF+g496mw+5DkQ8C3gC9V1c9nMUtVvVNVmxnsWW8Eru17m2dL8hngZFU9N+1tL+KWqvo4g6eZX0zy2ws/OaXHZazTti9nmrEfAy5fsLyxWzcrQ50Ke9KSXMQg9H+oqn+e5SwAVXUaeJrB4fLqJPP/v8Q0Hp9PAJ9N8hrwGIND+YdmMAdVdaz7eBL4NoNfgNN+XMY6bftyphn7s8A13Sutq4C7GZyOelamfirsJAEeAQ5V1ddmNUuSjyRZ3d3+AIPXDQ4xiP7Oac1RVfdX1caq2sTg5+H7VfX5ac+R5INJLpm/DXwKOMiUH5fq+7Ttfb/wcdYLDZ8GfsLg+eGfTXG73wSOA//N4LfnPQyeG+4DXgW+B6ydwhy3MDgE+yGD6+e92P03meoswMeAF7o5DgJ/3q2/CngGOAL8E/CrU3yMbgX2zGKObnsvdf9env/ZnNHPyGbgQPfYfAdYM6k5fAed1AhfoJMaYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ij/ARy0X2QY9RxEAAAAAElFTkSuQmCC",
       "text/plain": [
@@ -224,10 +215,10 @@
      },
      "metadata": {
       "needs_background": "light"
-     }
+     },
+     "output_type": "display_data"
     },
     {
-     "output_type": "display_data",
      "data": {
       "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAMqUlEQVR4nO3dUawc5XnG8f9TG5c0obENqWVhqI1AQVwEE1kUFFQRV0RuGgVfIESUSm6Fem5SiaiVEmiltqlUqdyEcFFVsoDGF23AJU1s+aLEcYjaK4MNprFxHJwUhC2DWwFK0gtUw9uLnVMdrGOf9dnZ3WN//590tDPfzu683t1n55vZ8TepKiRd+n5l2gVImgzDLjXCsEuNMOxSIwy71AjDLjVipLAn2ZLkWJLjSR7sqyhJ/ctif2dPsgz4CXAXcAJ4HvhCVb3cX3mS+rJ8hMfeChyvqp8BJHkSuBs4Z9iTeAaPNGZVlfnaR+nGXw28Pmf+RNcmaQkaZcs+lCQzwMy41yPp/EYJ+0ngmjnz67q2D6iq7cB2sBsvTdMo3fjngRuSbEiyArgP2N1PWZL6tugte1WdSfLHwDPAMuCJqjrSW2WSerXon94WtTK78dLYjeNovKSLiGGXGmHYpUYYdqkRhl1qhGGXGmHYpUYYdqkRhl1qhGGXGmHYpUYYdqkRhl1qhGGXGmHYpUYYdqkRhl1qhGGXGmHYpUYYdqkRhl1qhGGXGmHYpUYYdqkRhl1qxIJhT/JEktNJDs9pW51kb5JXuttV4y1T0qiG2bJ/E9hyVtuDwL6qugHY181LWsIWDHtV/Rvw1lnNdwM7uukdwNZ+y5LUt8Xus6+pqlPd9BvAmp7qkTQmi75k86yqqvNdnTXJDDAz6nokjWaxW/Y3k6wF6G5Pn2vBqtpeVZuqatMi1yWpB4sN+25gWze9DdjVTzmSxiVV5+yBDxZIvgXcCVwFvAn8JfBdYCdwLfAacG9VnX0Qb77nOv/KJI2sqjJf+4Jh75Nhl8bvXGH3DDqpEYZdaoRhlxph2KVGGHapEYZdaoRhlxox8rnxLZnkOQmarGTen6YvKW7ZpUYYdqkRhl1qhGGXGmHYpUYYdqkRhl1qhGGXGmHYpUYYdqkRhl1qhGGXGmHYpUYYdqkRhl1qhGGXGmHYpUYsGPYk1yR5NsnLSY4keaBrX51kb5JXuttV4y9X0mINc623tcDaqnohyRXAQWAr8AfAW1X1t0keBFZV1VcXeK6Lelwnh6W6dF1Kw1It+vJPVXWqql7opn8BHAWuBu4GdnSL7WDwBSBpibqgffYk64FbgP3Amqo61d31BrCm39Ik9Wno0WWTfAT4NvDlqvr53G5PVdW5uuhJZoCZUQuVNJqhLtmc5DJgD/BMVX29azsG3FlVp7r9+h9W1ccXeJ6LeqfXffZLl/vsQAavwuPA0dmgd3YD27rpbcCuUYuUND7DHI2/A/h34EfA+13znzHYb98JXAu8BtxbVW8t8FwX9abRLfulq4Ut+1Dd+L4Ydi1VLYTdM+ikRhh2qRGGXWqEYZcaYdilRhh2qRGGXWqEYZcaYdilRhh2qRGGXWqEYZcaYdilRhh2qRGGXWqEYZcaYdilRhh2qRGGXWqEYZcaYdilRhh2qRGGXWqEYZcaYdilRgxzrbfLkzyX5KUkR5J8rWvfkGR/kuNJnkqyYvzlSlqsYbbs7wKbq+pmYCOwJcltwMPAI1V1PfA2cP/YqpQ0sgXDXgO/7GYv6/4K2Aw83bXvALaOo0BJ/Rhqnz3JsiSHgNPAXuCnwDtVdaZb5ARw9VgqlNSLocJeVe9V1UZgHXArcOOwK0gyk+RAkgOLK1FSHy7oaHxVvQM8C9wOrEyyvLtrHXDyHI/ZXlWbqmrTKIVKGs0wR+M/lmRlN/0h4C7gKIPQ39Mttg3YNaYaJfUgVXX+BZJPMDgAt4zBl8POqvrrJNcBTwKrgReB36+qdxd4rvOvbIlb6LXSxSvJtEvoTVXN+49ZMOx9MuxaqloIu2fQSY0w7FIjDLvUCMMuNcKwS40w7FIjDLvUCMMuNcKwS40w7FIjDLvUCMMuNcKwS40w7FIjDLvUCMMuNcKwS40w7FIjDLvUCMMuNcKwS40w7FIjDLvUCMMuNcKwS40YOuzdZZtfTLKnm9+QZH+S40meSrJifGVKGtWFbNkfYHBBx1kPA49U1fXA28D9fRYmqV9DhT3JOuD3gMe6+QCbgae7RXYAW8dQn6SeDLtl/wbwFeD9bv5K4J2qOtPNnwCu7rc0SX0a5vrsnwNOV9XBxawgyUySA0kOLObxkvqxfIhlPgV8PslngcuBXwceBVYmWd5t3dcBJ+d7cFVtB7bDxX/JZulituCWvaoeqqp1VbUeuA/4QVV9EXgWuKdbbBuwa2xVShrZKL+zfxX4kyTHGezDP95PSZLGIVWT61lf7N34Sb5WmqzBD0yXhqqa9x/jGXRSIwy71AjDLjXCsEuNMOxSIwy71AjDLjXCsEuNMOxSIwy71AjDLjXCsEuNMOxSIwy71AjDLjXCsEuNMOxSIwy71AjDLjXCsEuNMOxSIwy71AjDLjXCsEuNGOZab+pcShcSUHuGCnuSV4FfAO8BZ6pqU5LVwFPAeuBV4N6qens8ZUoa1YV04z9dVRuralM3/yCwr6puAPZ185KWqFH22e8GdnTTO4CtI1cjaWyGDXsB30tyMMlM17amqk51028Aa3qvTlJvhj1Ad0dVnUzyG8DeJD+ee2dV1bmu0Np9OczMd5+kybngSzYn+Svgl8AfAXdW1akka4EfVtXHF3is1zyWxmzRl2xO8uEkV8xOA58BDgO7gW3dYtuAXf2UKmkcFtyyJ7kO+E43uxz4p6r6myRXAjuBa4HXGPz09tYCz+WWXRqzc23ZL7gbPwrDLo3forvxki4Nhl1qhGGXGmHYpUYYdqkRhl1qhGGXGmHYpUYYdqkRhl1qhGGXGmHYpUYYdqkRhl1qhGGXGmHYpUYYdqkRhl1qhGGXGmHYpUYYdqkRhl1qhGGXGmHYpUYYdqkRQ4U9ycokTyf5cZKjSW5PsjrJ3iSvdLerxl2spMUbdsv+KPCvVXUjcDNwFHgQ2FdVNwD7unlJS9QwF3b8KHAIuK7mLJzkGF6yWVpyRrnW2wbgv4B/SPJikse6SzevqapT3TJvAGv6KVXSOAwT9uXAJ4G/r6pbgP/hrC57t8Wfd6udZCbJgSQHRi1W0uINE/YTwImq2t/NP80g/G923Xe629PzPbiqtlfVpqra1EfBkhZnwbBX1RvA60lm98d/B3gZ2A1s69q2AbvGUqGkXix4gA4gyUbgMWAF8DPgDxl8UewErgVeA+6tqrcWeB4P0Eljdq4DdEOFvS+GXRq/UY7GS7oEGHapEYZdaoRhlxph2KVGGHapEYZdasTyCa/vvxmcgHNVNz1NS6EGsI6zWccHXWgdv3muOyZ6Us3/rzQ5MO1z5ZdCDdZhHZOsw2681AjDLjViWmHfPqX1zrUUagDrOJt1fFBvdUxln13S5NmNlxox0bAn2ZLkWJLjSSY2Gm2SJ5KcTnJ4TtvEh8JOck2SZ5O8nORIkgemUUuSy5M8l+Slro6vde0bkuzv3p+nkqwYZx1z6lnWjW+4Z1p1JHk1yY+SHJodQm1Kn5GxDds+sbAnWQb8HfC7wE3AF5LcNKHVfxPYclbbNIbCPgP8aVXdBNwGfKl7DSZdy7vA5qq6GdgIbElyG/Aw8EhVXQ+8Ddw/5jpmPcBgePJZ06rj01W1cc5PXdP4jIxv2PaqmsgfcDvwzJz5h4CHJrj+9cDhOfPHgLXd9Frg2KRqmVPDLuCuadYC/BrwAvBbDE7eWD7f+zXG9a/rPsCbgT1AplTHq8BVZ7VN9H0BPgr8J92xtL7rmGQ3/mrg9TnzJ7q2aZnqUNhJ1gO3APunUUvXdT7EYKDQvcBPgXeq6ky3yKTen28AXwHe7+avnFIdBXwvycEkM13bpN+XsQ7b7gE6zj8U9jgk+QjwbeDLVfXzadRSVe9V1UYGW9ZbgRvHvc6zJfkccLqqDk563fO4o6o+yWA380tJfnvunRN6X0Yatn0hkwz7SeCaOfPrurZpGWoo7L4luYxB0P+xqv5lmrUAVNU7wLMMussrk8z+f4lJvD+fAj6f5FXgSQZd+UenUAdVdbK7PQ18h8EX4KTfl5GGbV/IJMP+PHBDd6R1BXAfg+Gop2XiQ2EnCfA4cLSqvj6tWpJ8LMnKbvpDDI4bHGUQ+nsmVUdVPVRV66pqPYPPww+q6ouTriPJh5NcMTsNfAY4zITflxr3sO3jPvBx1oGGzwI/YbB/+OcTXO+3gFPA/zL49ryfwb7hPuAV4PvA6gnUcQeDLth/MLh+3qHuNZloLcAngBe7Og4Df9G1Xwc8BxwH/hn41Qm+R3cCe6ZRR7e+l7q/I7OfzSl9RjYCB7r35rvAqr7q8Aw6qREeoJMaYdilRhh2qRGGXWqEYZcaYdilRhh2qRGGXWrE/wE89HdrXg0aagAAAABJRU5ErkJggg==",
       "text/plain": [
@@ -236,36 +227,35 @@
      },
      "metadata": {
       "needs_background": "light"
-     }
+     },
+     "output_type": "display_data"
     }
    ],
-   "metadata": {}
+   "source": [
+    "test=binary_generator(2,(img_size,img_size,1),frac=frac_value)\n",
+    "imgs, y=next(test)\n",
+    "\n",
+    "display_img(imgs[0])\n",
+    "display_img(imgs[1])\n",
+    "print(\"Norm L2 \"+str(np.linalg.norm(imgs[1])))\n",
+    "print(\"Norm L2(count pixels) \"+str(math.sqrt(np.size(imgs[1][imgs[1]==1]))))"
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "#### for ternary generator"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 6,
-   "source": [
-    "test=ternary_generator(4,(img_size,img_size,1),frac=frac_value)\n",
-    "imgs, y=next(test)\n",
-    "\n",
-    "for i in range(4):\n",
-    "    display_img(0.5*(imgs[i]+1.0)) # we ensure that there is no negative value wehn displaying images\n",
-    "\n",
-    "print(\"Norm L2(imgs[2]-imgs[0])\"+str(np.linalg.norm(imgs[2]-imgs[0])))\n",
-    "print(\"Norm L2(imgs[2]) \"+str(np.linalg.norm(imgs[2])))\n",
-    "print(\"Norm L2(count pixels) \"+str(math.sqrt(np.size(imgs[2][imgs[2]==-1]))))"
-   ],
+   "metadata": {},
    "outputs": [
     {
-     "output_type": "stream",
      "name": "stdout",
+     "output_type": "stream",
      "text": [
       "Norm L2(imgs[2]-imgs[0])35.0\n",
       "Norm L2(imgs[2]) 35.0\n",
@@ -273,7 +263,6 @@
      ]
     },
     {
-     "output_type": "display_data",
      "data": {
       "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAMvklEQVR4nO3df6jd9X3H8edric7etvNq60I0smQYKv4xY7lYRSmrzpK5UvOHiFJGGIH7jxuWFTrdYFDYH/WfWv8Yg0t1zR+u6mxdRErbLLWMwYheq7bR1Jo6xYRous3Qbhe6xb73x/mmXMON9+Se7zkn2+f5gHDO93u+x+8bz33e8+Mevt9UFZL+//u1aQ8gaTKMXWqEsUuNMHapEcYuNcLYpUaMFHuS7UleTnIoyd19DSWpf1nr39mTrAN+DNwEHAaeAe6oqpf6G09SX9aPcN+rgUNV9SpAkoeBW4DTxj4zM1Ozs7Mj7FLSezl+/DhLS0tZ6bZRYr8EeGPZ8mHgY+91h9nZWebn50fYpaT3srCwcNrbxv4BXZL5JItJFpeWlsa9O0mnMUrsR4BLly1v6ta9S1UtVNVcVc3NzMyMsDtJoxgl9meArUm2JDkXuB14op+xJPVtze/Zq+pEkj8Gvg2sAx6sqhd7m0xSr0b5gI6q+ibwzZ5mkTRGfoNOaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdasSqsSd5MMmxJAeWrbswyd4kr3SXF4x3TEmjGuaZ/avA9lPW3Q3sq6qtwL5uWdJZbNXYq+qfgP84ZfUtwO7u+m5gR79jSerbWt+zb6iqo931N4ENPc0jaUxG/oCuqgqo092eZD7JYpLFpaWlUXcnaY3WGvtbSTYCdJfHTrdhVS1U1VxVzc3MzKxxd5JGtdbYnwB2dtd3Anv6GUfSuAzzp7evAf8CfCTJ4SS7gC8CNyV5Bfi9blnSWWz9ahtU1R2nuenGnmeRNEZ+g05qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qxDCnf7o0yVNJXkryYpK7uvUXJtmb5JXu8oLxjytprYZ5Zj8BfK6qrgCuAe5McgVwN7CvqrYC+7plSWepVWOvqqNV9f3u+s+Bg8AlwC3A7m6z3cCOMc0oqQdn9J49yWbgKmA/sKGqjnY3vQls6Hc0SX0aOvYkHwC+Dny2qn62/LaqKqBOc7/5JItJFpeWlkYaVtLaDRV7knMYhP5QVX2jW/1Wko3d7RuBYyvdt6oWqmququZmZmb6mFnSGgzzaXyAB4CDVfWlZTc9Aezsru8E9vQ/nqS+rB9im+uAPwR+mOT5bt2fA18EHk2yC3gduG0sE0rqxaqxV9U/AznNzTf2O46kcfEbdFIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjhjnX23lJnk7yQpIXk3yhW78lyf4kh5I8kuTc8Y8raa2GeWb/BXBDVV0JbAO2J7kGuBe4r6ouA94Gdo1tSkkjWzX2GvjPbvGc7l8BNwCPdet3AzvGMaCkfgx7fvZ13RlcjwF7gZ8Ax6vqRLfJYeCSsUwoqRdDxV5V71TVNmATcDVw+bA7SDKfZDHJ4tLS0tqmlDSyM/o0vqqOA08B1wKzSU6e8nkTcOQ091moqrmqmpuZmRllVkkjGObT+IuSzHbX3wfcBBxkEP2t3WY7gT1jmlFSD9avvgkbgd1J1jH45fBoVT2Z5CXg4SR/BTwHPDDGOSWNaNXYq+oHwFUrrH+Vwft3Sf8H+A06qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRFDx96dtvm5JE92y1uS7E9yKMkjSc4d35iSRnUmz+x3MTih40n3AvdV1WXA28CuPgeT1K+hYk+yCfgD4CvdcoAbgMe6TXYDO8Ywn6SeDPvM/mXg88Avu+UPAcer6kS3fBi4pN/RJPVpmPOzfwo4VlXPrmUHSeaTLCZZXFpaWst/QlIPhjk/+3XAp5PcDJwH/AZwPzCbZH337L4JOLLSnatqAVgAuPjii6uXqSWdsVWf2avqnqraVFWbgduB71bVZ4CngFu7zXYCe8Y2paSRjfJ39j8D/jTJIQbv4R/oZyRJ4zDMy/hfqarvAd/rrr8KXN3/SJLGwW/QSY0wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40Y6owwSV4Dfg68A5yoqrkkFwKPAJuB14Dbqurt8YwpaVRn8sz+iaraVlVz3fLdwL6q2grs65YlnaVGeRl/C7C7u74b2DHyNJLGZtjYC/hOkmeTzHfrNlTV0e76m8CG3qeT1Jthz+J6fVUdSfKbwN4kP1p+Y1VVklrpjt0vh3mA888/f6RhJa3dUM/sVXWkuzwGPM7gVM1vJdkI0F0eO819F6pqrqrmZmZm+pla0hlbNfYk70/ywZPXgU8CB4AngJ3dZjuBPeMaUtLohnkZvwF4PMnJ7f+uqr6V5Bng0SS7gNeB28Y3pqRRrRp7Vb0KXLnC+n8HbhzHUJL65zfopEYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYMFXuS2SSPJflRkoNJrk1yYZK9SV7pLi8Y97CS1m7YZ/b7gW9V1eUMTgV1ELgb2FdVW4F93bKks9QwZ3E9H/g48ABAVf13VR0HbgF2d5vtBnaMZ0RJfRjmmX0L8FPgb5M8l+Qr3ambN1TV0W6bNxmc7VXSWWqY2NcDHwX+pqquAv6LU16yV1UBtdKdk8wnWUyyuLS0NOq8ktZomNgPA4eran+3/BiD+N9KshGguzy20p2raqGq5qpqbmZmpo+ZJa3BqrFX1ZvAG0k+0q26EXgJeALY2a3bCewZy4SSerF+yO3+BHgoybnAq8AfMfhF8WiSXcDrwG3jGVFSH4aKvaqeB+ZWuOnGXqeRNDZ+g05qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGpHB19ontLPkpwy+gPNh4N8mtuOVnQ0zgHOcyjne7Uzn+K2qumilGyYa+692mixW1Upf0mlqBudwjknO4ct4qRHGLjViWrEvTGm/y50NM4BznMo53q23Oabynl3S5PkyXmrERGNPsj3Jy0kOJZnY0WiTPJjkWJIDy9ZN/FDYSS5N8lSSl5K8mOSuacyS5LwkTyd5oZvjC936LUn2d4/PI93xC8Yuybru+IZPTmuOJK8l+WGS55Msduum8TMytsO2Tyz2JOuAvwZ+H7gCuCPJFRPa/VeB7aesm8ahsE8An6uqK4BrgDu7/weTnuUXwA1VdSWwDdie5BrgXuC+qroMeBvYNeY5TrqLweHJT5rWHJ+oqm3L/tQ1jZ+R8R22vaom8g+4Fvj2suV7gHsmuP/NwIFlyy8DG7vrG4GXJzXLshn2ADdNcxZgBvg+8DEGX95Yv9LjNcb9b+p+gG8AngQypTleAz58yrqJPi7A+cC/0n2W1vcck3wZfwnwxrLlw926aZnqobCTbAauAvZPY5bupfPzDA4Uuhf4CXC8qk50m0zq8fky8Hngl93yh6Y0RwHfSfJskvlu3aQfl7Eett0P6HjvQ2GPQ5IPAF8HPltVP5vGLFX1TlVtY/DMejVw+bj3eaoknwKOVdWzk973Cq6vqo8yeJt5Z5KPL79xQo/LSIdtX80kYz8CXLpseVO3blqGOhR235KcwyD0h6rqG9OcBaAGZ/d5isHL5dkkJ49LOInH5zrg00leAx5m8FL+/inMQVUd6S6PAY8z+AU46cdlpMO2r2aSsT8DbO0+aT0XuJ3B4ainZeKHwk4SBqfROlhVX5rWLEkuSjLbXX8fg88NDjKI/tZJzVFV91TVpqrazODn4btV9ZlJz5Hk/Uk+ePI68EngABN+XGrch20f9wcfp3zQcDPwYwbvD/9igvv9GnAU+B8Gvz13MXhvuA94BfhH4MIJzHE9g5dgPwCe7/7dPOlZgN8BnuvmOAD8Zbf+t4GngUPA3wO/PsHH6HeBJ6cxR7e/F7p/L5782ZzSz8g2YLF7bP4BuKCvOfwGndQIP6CTGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9SI/wX2pja/9OIdNQAAAABJRU5ErkJggg==",
       "text/plain": [
@@ -282,10 +271,10 @@
      },
      "metadata": {
       "needs_background": "light"
-     }
+     },
+     "output_type": "display_data"
     },
     {
-     "output_type": "display_data",
      "data": {
       "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAMvklEQVR4nO3df6jd9X3H8edric7etvNq60I0smQYKv4xY7lYRSmrzpK5UvOHiFJGGIH7jxuWFTrdYFDYH/WfWv8Yg0t1zR+u6mxdRErbLLWMwYheq7bR1Jo6xYRous3Qbhe6xb73x/mmXMON9+Se7zkn2+f5gHDO93u+x+8bz33e8+Mevt9UFZL+//u1aQ8gaTKMXWqEsUuNMHapEcYuNcLYpUaMFHuS7UleTnIoyd19DSWpf1nr39mTrAN+DNwEHAaeAe6oqpf6G09SX9aPcN+rgUNV9SpAkoeBW4DTxj4zM1Ozs7Mj7FLSezl+/DhLS0tZ6bZRYr8EeGPZ8mHgY+91h9nZWebn50fYpaT3srCwcNrbxv4BXZL5JItJFpeWlsa9O0mnMUrsR4BLly1v6ta9S1UtVNVcVc3NzMyMsDtJoxgl9meArUm2JDkXuB14op+xJPVtze/Zq+pEkj8Gvg2sAx6sqhd7m0xSr0b5gI6q+ibwzZ5mkTRGfoNOaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdasSqsSd5MMmxJAeWrbswyd4kr3SXF4x3TEmjGuaZ/avA9lPW3Q3sq6qtwL5uWdJZbNXYq+qfgP84ZfUtwO7u+m5gR79jSerbWt+zb6iqo931N4ENPc0jaUxG/oCuqgqo092eZD7JYpLFpaWlUXcnaY3WGvtbSTYCdJfHTrdhVS1U1VxVzc3MzKxxd5JGtdbYnwB2dtd3Anv6GUfSuAzzp7evAf8CfCTJ4SS7gC8CNyV5Bfi9blnSWWz9ahtU1R2nuenGnmeRNEZ+g05qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qxDCnf7o0yVNJXkryYpK7uvUXJtmb5JXu8oLxjytprYZ5Zj8BfK6qrgCuAe5McgVwN7CvqrYC+7plSWepVWOvqqNV9f3u+s+Bg8AlwC3A7m6z3cCOMc0oqQdn9J49yWbgKmA/sKGqjnY3vQls6Hc0SX0aOvYkHwC+Dny2qn62/LaqKqBOc7/5JItJFpeWlkYaVtLaDRV7knMYhP5QVX2jW/1Wko3d7RuBYyvdt6oWqmququZmZmb6mFnSGgzzaXyAB4CDVfWlZTc9Aezsru8E9vQ/nqS+rB9im+uAPwR+mOT5bt2fA18EHk2yC3gduG0sE0rqxaqxV9U/AznNzTf2O46kcfEbdFIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjhjnX23lJnk7yQpIXk3yhW78lyf4kh5I8kuTc8Y8raa2GeWb/BXBDVV0JbAO2J7kGuBe4r6ouA94Gdo1tSkkjWzX2GvjPbvGc7l8BNwCPdet3AzvGMaCkfgx7fvZ13RlcjwF7gZ8Ax6vqRLfJYeCSsUwoqRdDxV5V71TVNmATcDVw+bA7SDKfZDHJ4tLS0tqmlDSyM/o0vqqOA08B1wKzSU6e8nkTcOQ091moqrmqmpuZmRllVkkjGObT+IuSzHbX3wfcBBxkEP2t3WY7gT1jmlFSD9avvgkbgd1J1jH45fBoVT2Z5CXg4SR/BTwHPDDGOSWNaNXYq+oHwFUrrH+Vwft3Sf8H+A06qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRFDx96dtvm5JE92y1uS7E9yKMkjSc4d35iSRnUmz+x3MTih40n3AvdV1WXA28CuPgeT1K+hYk+yCfgD4CvdcoAbgMe6TXYDO8Ywn6SeDPvM/mXg88Avu+UPAcer6kS3fBi4pN/RJPVpmPOzfwo4VlXPrmUHSeaTLCZZXFpaWst/QlIPhjk/+3XAp5PcDJwH/AZwPzCbZH337L4JOLLSnatqAVgAuPjii6uXqSWdsVWf2avqnqraVFWbgduB71bVZ4CngFu7zXYCe8Y2paSRjfJ39j8D/jTJIQbv4R/oZyRJ4zDMy/hfqarvAd/rrr8KXN3/SJLGwW/QSY0wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40Y6owwSV4Dfg68A5yoqrkkFwKPAJuB14Dbqurt8YwpaVRn8sz+iaraVlVz3fLdwL6q2grs65YlnaVGeRl/C7C7u74b2DHyNJLGZtjYC/hOkmeTzHfrNlTV0e76m8CG3qeT1Jthz+J6fVUdSfKbwN4kP1p+Y1VVklrpjt0vh3mA888/f6RhJa3dUM/sVXWkuzwGPM7gVM1vJdkI0F0eO819F6pqrqrmZmZm+pla0hlbNfYk70/ywZPXgU8CB4AngJ3dZjuBPeMaUtLohnkZvwF4PMnJ7f+uqr6V5Bng0SS7gNeB28Y3pqRRrRp7Vb0KXLnC+n8HbhzHUJL65zfopEYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYMFXuS2SSPJflRkoNJrk1yYZK9SV7pLi8Y97CS1m7YZ/b7gW9V1eUMTgV1ELgb2FdVW4F93bKks9QwZ3E9H/g48ABAVf13VR0HbgF2d5vtBnaMZ0RJfRjmmX0L8FPgb5M8l+Qr3ambN1TV0W6bNxmc7VXSWWqY2NcDHwX+pqquAv6LU16yV1UBtdKdk8wnWUyyuLS0NOq8ktZomNgPA4eran+3/BiD+N9KshGguzy20p2raqGq5qpqbmZmpo+ZJa3BqrFX1ZvAG0k+0q26EXgJeALY2a3bCewZy4SSerF+yO3+BHgoybnAq8AfMfhF8WiSXcDrwG3jGVFSH4aKvaqeB+ZWuOnGXqeRNDZ+g05qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGpHB19ontLPkpwy+gPNh4N8mtuOVnQ0zgHOcyjne7Uzn+K2qumilGyYa+692mixW1Upf0mlqBudwjknO4ct4qRHGLjViWrEvTGm/y50NM4BznMo53q23Oabynl3S5PkyXmrERGNPsj3Jy0kOJZnY0WiTPJjkWJIDy9ZN/FDYSS5N8lSSl5K8mOSuacyS5LwkTyd5oZvjC936LUn2d4/PI93xC8Yuybru+IZPTmuOJK8l+WGS55Msduum8TMytsO2Tyz2JOuAvwZ+H7gCuCPJFRPa/VeB7aesm8ahsE8An6uqK4BrgDu7/weTnuUXwA1VdSWwDdie5BrgXuC+qroMeBvYNeY5TrqLweHJT5rWHJ+oqm3L/tQ1jZ+R8R22vaom8g+4Fvj2suV7gHsmuP/NwIFlyy8DG7vrG4GXJzXLshn2ADdNcxZgBvg+8DEGX95Yv9LjNcb9b+p+gG8AngQypTleAz58yrqJPi7A+cC/0n2W1vcck3wZfwnwxrLlw926aZnqobCTbAauAvZPY5bupfPzDA4Uuhf4CXC8qk50m0zq8fky8Hngl93yh6Y0RwHfSfJskvlu3aQfl7Eett0P6HjvQ2GPQ5IPAF8HPltVP5vGLFX1TlVtY/DMejVw+bj3eaoknwKOVdWzk973Cq6vqo8yeJt5Z5KPL79xQo/LSIdtX80kYz8CXLpseVO3blqGOhR235KcwyD0h6rqG9OcBaAGZ/d5isHL5dkkJ49LOInH5zrg00leAx5m8FL+/inMQVUd6S6PAY8z+AU46cdlpMO2r2aSsT8DbO0+aT0XuJ3B4ainZeKHwk4SBqfROlhVX5rWLEkuSjLbXX8fg88NDjKI/tZJzVFV91TVpqrazODn4btV9ZlJz5Hk/Uk+ePI68EngABN+XGrch20f9wcfp3zQcDPwYwbvD/9igvv9GnAU+B8Gvz13MXhvuA94BfhH4MIJzHE9g5dgPwCe7/7dPOlZgN8BnuvmOAD8Zbf+t4GngUPA3wO/PsHH6HeBJ6cxR7e/F7p/L5782ZzSz8g2YLF7bP4BuKCvOfwGndQIP6CTGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9SI/wX2pja/9OIdNQAAAABJRU5ErkJggg==",
       "text/plain": [
@@ -294,10 +283,10 @@
      },
      "metadata": {
       "needs_background": "light"
-     }
+     },
+     "output_type": "display_data"
     },
     {
-     "output_type": "display_data",
      "data": {
       "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAM20lEQVR4nO3df6jd9X3H8edric7etvNq60I0sjgMFf+YsVysopRVZ8lcqflDRCkjjMD9xw3LCp1usCDsj/pPrX+MQaiu+cNVna2LSGmbppYxGNFr1TaaWlOnmBBNtxna7UK32Pf+ON9013DjPbnne87J9nk+IJzz/Z7v8fvGc5/3/LiH7zdVhaT//35t2gNImgxjlxph7FIjjF1qhLFLjTB2qREjxZ5kS5KXkxxMcldfQ0nqX1b7d/Yka4AfAzcCh4BngNur6qX+xpPUl7Uj3Pcq4GBVvQqQ5GHgZuCUsc/MzNTs7OwIu5T0Xo4dO8bi4mKWu22U2C8C3liyfAj42HvdYXZ2lvn5+RF2Kem97Ny585S3jf0DuiTzSRaSLCwuLo57d5JOYZTYDwMXL1ne0K17l6raWVVzVTU3MzMzwu4kjWKU2J8BNiW5JMnZwG3AE/2MJalvq37PXlXHk/wx8C1gDfBgVb3Y22SSejXKB3RU1TeAb/Q0i6Qx8ht0UiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiNWjD3Jg0mOJtm/ZN35SfYkeaW7PG+8Y0oa1TDP7F8Btpy07i5gb1VtAvZ2y5LOYCvGXlX/CPz7SatvBnZ113cBW/sdS1LfVvuefV1VHemuvwms62keSWMy8gd0VVVAner2JPNJFpIsLC4ujro7Sau02tjfSrIeoLs8eqoNq2pnVc1V1dzMzMwqdydpVKuN/QlgW3d9G7C7n3Ekjcswf3r7KvDPwEeSHEqyHfgCcGOSV4Df65YlncHWrrRBVd1+iptu6HkWSWPkN+ikRhi71Ahjlxph7FIjjF1qhLFLjTB2qREr/p1d/+uee+6Z9ggakx07dkx7hLHzmV1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdasQwp3+6OMlTSV5K8mKSO7v15yfZk+SV7vK88Y8rabWGeWY/Dnyuqi4HrgbuSHI5cBewt6o2AXu7ZUlnqBVjr6ojVfX97vrPgQPARcDNwK5us13A1jHNKKkHp/WePclG4EpgH7Cuqo50N70JrOt3NEl9Gjr2JB8AvgZ8tqp+tvS2qiqgTnG/+SQLSRYWFxdHGlbS6g0Ve5KzGIT+UFV9vVv9VpL13e3rgaPL3beqdlbVXFXNzczM9DGzpFUY5tP4AA8AB6rqi0tuegLY1l3fBuzufzxJfRnmJBHXAn8I/DDJ8926Pwe+ADyaZDvwOnDrWCaU1IsVY6+qfwJyiptv6HccSePiN+ikRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRgxzrrdzkjyd5IUkLya5p1t/SZJ9SQ4meSTJ2eMfV9JqDfPM/gvg+qq6AtgMbElyNXAvcF9VXQq8DWwf25SSRrZi7DXwH93iWd2/Aq4HHuvW7wK2jmNASf0Y9vzsa7ozuB4F9gA/AY5V1fFuk0PARWOZUFIvhoq9qt6pqs3ABuAq4LJhd5BkPslCkoXFxcXVTSlpZKf1aXxVHQOeAq4BZpOcOOXzBuDwKe6zs6rmqmpuZmZmlFkljWCYT+MvSDLbXX8fcCNwgEH0t3SbbQN2j2lGST1Yu/ImrAd2JVnD4JfDo1X1ZJKXgIeT/BXwHPDAGOeUNKIVY6+qHwBXLrP+VQbv3yX9H+A36KRGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGDB17d9rm55I82S1fkmRfkoNJHkly9vjGlDSq03lmv5PBCR1PuBe4r6ouBd4Gtvc5mKR+DRV7kg3AHwBf7pYDXA881m2yC9g6hvkk9WTYZ/YvAZ8Hftktfwg4VlXHu+VDwEX9jiapT8Ocn/1TwNGqenY1O0gyn2QhycLi4uJq/hOSejDM+dmvBT6d5CbgHOA3gPuB2SRru2f3DcDh5e5cVTuBnQAXXnhh9TK1pNO24jN7Vd1dVRuqaiNwG/DdqvoM8BRwS7fZNmD32KaUNLJR/s7+Z8CfJjnI4D38A/2MJGkchnkZ/ytV9T3ge931V4Gr+h9J0jj4DTqpEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNOK2TRLRux44d0x5BWrWhYk/yGvBz4B3geFXNJTkfeATYCLwG3FpVb49nTEmjOp2X8Z+oqs1VNdct3wXsrapNwN5uWdIZapT37DcDu7rru4CtI08jaWyGjb2Abyd5Nsl8t25dVR3prr8JrOt9Okm9GfYDuuuq6nCS3wT2JPnR0hurqpLUcnfsfjnMA5x77rkjDStp9YZ6Zq+qw93lUeBxBqdqfivJeoDu8ugp7ruzquaqam5mZqafqSWdthVjT/L+JB88cR34JLAfeALY1m22Ddg9riEljW6Yl/HrgMeTnNj+76rqm0meAR5Nsh14Hbh1fGNKGtWKsVfVq8AVy6z/N+CGcQwlqX9+XVZqhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qxFCxJ5lN8liSHyU5kOSaJOcn2ZPkle7yvHEPK2n1hn1mvx/4ZlVdxuBUUAeAu4C9VbUJ2NstSzpDDXMW13OBjwMPAFTVf1XVMeBmYFe32S5g63hGlNSHYZ7ZLwF+CvxtkueSfLk7dfO6qjrSbfMmg7O9SjpDDRP7WuCjwN9U1ZXAf3LSS/aqKqCWu3OS+SQLSRYWFxdHnVfSKg0T+yHgUFXt65YfYxD/W0nWA3SXR5e7c1XtrKq5qpqbmZnpY2ZJq7Bi7FX1JvBGko90q24AXgKeALZ167YBu8cyoaRerB1yuz8BHkpyNvAq8EcMflE8mmQ78Dpw63hGlNSHoWKvqueBuWVuuqHXaSSNjd+gkxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGZPC19gntLPkpgy/gfBj414nteHlnwgzgHCdzjnc73Tl+q6ouWO6Gicb+q50mC1W13Jd0mprBOZxjknP4Ml5qhLFLjZhW7DuntN+lzoQZwDlO5hzv1tscU3nPLmnyfBkvNWKisSfZkuTlJAeTTOxotEkeTHI0yf4l6yZ+KOwkFyd5KslLSV5Mcuc0ZklyTpKnk7zQzXFPt/6SJPu6x+eR7vgFY5dkTXd8wyenNUeS15L8MMnzSRa6ddP4GRnbYdsnFnuSNcBfA78PXA7cnuTyCe3+K8CWk9ZN41DYx4HPVdXlwNXAHd3/g0nP8gvg+qq6AtgMbElyNXAvcF9VXQq8DWwf8xwn3Mng8OQnTGuOT1TV5iV/6prGz8j4DtteVRP5B1wDfGvJ8t3A3RPc/0Zg/5Lll4H13fX1wMuTmmXJDLuBG6c5CzADfB/4GIMvb6xd7vEa4/43dD/A1wNPApnSHK8BHz5p3UQfF+Bc4F/oPkvre45Jvoy/CHhjyfKhbt20TPVQ2Ek2AlcC+6YxS/fS+XkGBwrdA/wEOFZVx7tNJvX4fAn4PPDLbvlDU5qjgG8neTbJfLdu0o/LWA/b7gd0vPehsMchyQeArwGfraqfTWOWqnqnqjYzeGa9Crhs3Ps8WZJPAUer6tlJ73sZ11XVRxm8zbwjyceX3jihx2Wkw7avZJKxHwYuXrK8oVs3LUMdCrtvSc5iEPpDVfX1ac4CUIOz+zzF4OXybJITxyWcxONzLfDpJK8BDzN4KX//FOagqg53l0eBxxn8Apz04zLSYdtXMsnYnwE2dZ+0ng3cxuBw1NMy8UNhJwmD02gdqKovTmuWJBckme2uv4/B5wYHGER/y6TmqKq7q2pDVW1k8PPw3ar6zKTnSPL+JB88cR34JLCfCT8uNe7Dto/7g4+TPmi4Cfgxg/eHfzHB/X4VOAL8N4PfntsZvDfcC7wCfAc4fwJzXMfgJdgPgOe7fzdNehbgd4Dnujn2A3/Zrf9t4GngIPD3wK9P8DH6XeDJaczR7e+F7t+LJ342p/QzshlY6B6bfwDO62sOv0EnNcIP6KRGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNeJ/AMelO0A/m/h2AAAAAElFTkSuQmCC",
       "text/plain": [
@@ -306,10 +295,10 @@
      },
      "metadata": {
       "needs_background": "light"
-     }
+     },
+     "output_type": "display_data"
     },
     {
-     "output_type": "display_data",
      "data": {
       "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAM1klEQVR4nO3df6jd9X3H8edric7etvNq60I0Mh2Gin/MWC5WUcqqs2RdqflDRCkjjMD9xw3LCp1uMFLYH/WfWv8Yg0t1zR+u6mxdRErbNLWMwYheq7bR1Jo6xYRous3Qbhe6xb73x/mmu4Yb78k933NOts/zAeGc7/d8j983nvu858c9fL+pKiT9//dr0x5A0mQYu9QIY5caYexSI4xdaoSxS40YKfYkW5O8lORgkrv6GkpS/7LWv7MnWQf8GLgJOAQ8DdxeVS/2N56kvqwf4b5XAwer6hWAJA8BNwOnjH1mZqZmZ2dH2KWkd3Ps2DGWlpay0m2jxH4R8Pqy5UPAR97tDrOzs8zPz4+wS0nvZmFh4ZS3jf0DuiTzSRaTLC4tLY17d5JOYZTYDwMXL1ve1K17h6paqKq5qpqbmZkZYXeSRjFK7E8Dm5NcmuRs4Dbg8X7GktS3Nb9nr6rjSf4Y+BawDnigql7obTJJvRrlAzqq6hvAN3qaRdIY+Q06qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRGrxp7kgSRHk+xftu78JHuSvNxdnjfeMSWNaphn9q8AW09adxewt6o2A3u7ZUlnsFVjr6p/BP79pNU3A7u667uAbf2OJalva33PvqGqjnTX3wA29DSPpDEZ+QO6qiqgTnV7kvkki0kWl5aWRt2dpDVaa+xvJtkI0F0ePdWGVbVQVXNVNTczM7PG3Uka1VpjfxzY3l3fDuzuZxxJ4zLMn96+Cvwz8KEkh5LsAL4A3JTkZeD3umVJZ7D1q21QVbef4qYbe55F0hj5DTqpEcYuNcLYpUYYu9QIY5caYexSI4xdasSqf2fX/9q5c+e0R9CYtPDY+swuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjVimNM/XZzkySQvJnkhyZ3d+vOT7Enycnd53vjHlbRWwzyzHwc+W1VXANcAdyS5ArgL2FtVm4G93bKkM9SqsVfVkar6fnf958AB4CLgZmBXt9kuYNuYZpTUg9N6z57kEuAqYB+woaqOdDe9AWzodzRJfRo69iTvA74GfKaqfrb8tqoqoE5xv/kki0kWl5aWRhpW0toNFXuSsxiE/mBVfb1b/WaSjd3tG4GjK923qhaqaq6q5mZmZvqYWdIaDPNpfID7gQNV9cVlNz0ObO+ubwd29z+epL4Mc5KI64A/BH6Y5Llu3Z8DXwAeSbIDeA24dSwTSurFqrFX1T8BOcXNN/Y7jqRx8Rt0UiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOGOdfbOUmeSvJ8kheSfL5bf2mSfUkOJnk4ydnjH1fSWg3zzP4L4IaquhLYAmxNcg1wD3BvVV0GvAXsGNuUkka2auw18B/d4lndvwJuAB7t1u8Cto1jQEn9GPb87Ou6M7geBfYAPwGOVdXxbpNDwEVjmVBSL4aKvarerqotwCbgauDyYXeQZD7JYpLFpaWltU0paWSn9Wl8VR0DngSuBWaTnDjl8ybg8Cnus1BVc1U1NzMzM8qskkYwzKfxFySZ7a6/B7gJOMAg+lu6zbYDu8c0o6QerF99EzYCu5KsY/DL4ZGqeiLJi8BDSf4KeBa4f4xzShrRqrFX1Q+Aq1ZY/wqD9++S/g/wG3RSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4aOvTtt87NJnuiWL02yL8nBJA8nOXt8Y0oa1ek8s9/J4ISOJ9wD3FtVlwFvATv6HExSv4aKPckm4A+AL3fLAW4AHu022QVsG8N8knoy7DP7l4DPAb/slj8AHKuq493yIeCifkeT1Kdhzs/+SeBoVT2zlh0kmU+ymGRxaWlpLf8JST0Y5vzs1wGfSvIJ4BzgN4D7gNkk67tn903A4ZXuXFULwALAhRdeWL1MLem0rfrMXlV3V9WmqroEuA34blV9GngSuKXbbDuwe2xTShrZKH9n/zPgT5McZPAe/v5+RpI0DsO8jP+Vqvoe8L3u+ivA1f2PJGkc/Aad1Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRpzWSSJat3PnzmmPIK3ZULEneRX4OfA2cLyq5pKcDzwMXAK8CtxaVW+NZ0xJozqdl/Efq6otVTXXLd8F7K2qzcDeblnSGWqU9+w3A7u667uAbSNPI2lsho29gG8neSbJfLduQ1Ud6a6/AWzofTpJvRn2A7rrq+pwkt8E9iT50fIbq6qS1Ep37H45zAOce+65Iw0rae2GemavqsPd5VHgMQanan4zyUaA7vLoKe67UFVzVTU3MzPTz9SSTtuqsSd5b5L3n7gOfBzYDzwObO822w7sHteQkkY3zMv4DcBjSU5s/3dV9c0kTwOPJNkBvAbcOr4xJY1q1dir6hXgyhXW/xtw4ziGktQ/vy4rNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkRxi41wtilRhi71Ahjlxph7FIjjF1qhLFLjTB2qRHGLjXC2KVGGLvUCGOXGmHsUiOMXWqEsUuNMHapEcYuNWKo2JPMJnk0yY+SHEhybZLzk+xJ8nJ3ed64h5W0dsM+s98HfLOqLmdwKqgDwF3A3qraDOztliWdoYY5i+u5wEeB+wGq6r+q6hhwM7Cr22wXsG08I0rqwzDP7JcCPwX+NsmzSb7cnbp5Q1Ud6bZ5g8HZXiWdoYaJfT3wYeBvquoq4D856SV7VRVQK905yXySxSSLS0tLo84raY2Gif0QcKiq9nXLjzKI/80kGwG6y6Mr3bmqFqpqrqrmZmZm+phZ0hqsGntVvQG8nuRD3aobgReBx4Ht3brtwO6xTCipF+uH3O5PgAeTnA28AvwRg18UjyTZAbwG3DqeESX1YajYq+o5YG6Fm27sdRpJY+M36KRGGLvUCGOXGmHsUiOMXWqEsUuNMHapERl8rX1CO0t+yuALOB8E/nViO17ZmTADOMfJnOOdTneO36qqC1a6YaKx/2qnyWJVrfQlnaZmcA7nmOQcvoyXGmHsUiOmFfvClPa73JkwAzjHyZzjnXqbYyrv2SVNni/jpUZMNPYkW5O8lORgkokdjTbJA0mOJtm/bN3ED4Wd5OIkTyZ5MckLSe6cxixJzknyVJLnuzk+362/NMm+7vF5uDt+wdglWdcd3/CJac2R5NUkP0zyXJLFbt00fkbGdtj2icWeZB3w18DvA1cAtye5YkK7/wqw9aR10zgU9nHgs1V1BXANcEf3/2DSs/wCuKGqrgS2AFuTXAPcA9xbVZcBbwE7xjzHCXcyODz5CdOa42NVtWXZn7qm8TMyvsO2V9VE/gHXAt9atnw3cPcE938JsH/Z8kvAxu76RuClSc2ybIbdwE3TnAWYAb4PfITBlzfWr/R4jXH/m7of4BuAJ4BMaY5XgQ+etG6ijwtwLvAvdJ+l9T3HJF/GXwS8vmz5ULduWqZ6KOwklwBXAfumMUv30vk5BgcK3QP8BDhWVce7TSb1+HwJ+Bzwy275A1Oao4BvJ3kmyXy3btKPy1gP2+4HdLz7obDHIcn7gK8Bn6mqn01jlqp6u6q2MHhmvRq4fNz7PFmSTwJHq+qZSe97BddX1YcZvM28I8lHl984ocdlpMO2r2aSsR8GLl62vKlbNy1DHQq7b0nOYhD6g1X19WnOAlCDs/s8yeDl8mySE8clnMTjcx3wqSSvAg8xeCl/3xTmoKoOd5dHgccY/AKc9OMy0mHbVzPJ2J8GNneftJ4N3MbgcNTTMvFDYScJg9NoHaiqL05rliQXJJntrr+HwecGBxhEf8uk5qiqu6tqU1VdwuDn4btV9elJz5HkvUnef+I68HFgPxN+XGrch20f9wcfJ33Q8AngxwzeH/7FBPf7VeAI8N8MfnvuYPDecC/wMvAd4PwJzHE9g5dgPwCe6/59YtKzAL8DPNvNsR/4y279bwNPAQeBvwd+fYKP0e8CT0xjjm5/z3f/Xjjxszmln5EtwGL32PwDcF5fc/gNOqkRfkAnNcLYpUYYu9QIY5caYexSI4xdaoSxS40wdqkR/wMyEDtD4uKLKQAAAABJRU5ErkJggg==",
       "text/plain": [
@@ -318,13 +307,25 @@
      },
      "metadata": {
       "needs_background": "light"
-     }
+     },
+     "output_type": "display_data"
     }
    ],
-   "metadata": {}
+   "source": [
+    "test=ternary_generator(4,(img_size,img_size,1),frac=frac_value)\n",
+    "imgs, y=next(test)\n",
+    "\n",
+    "for i in range(4):\n",
+    "    display_img(0.5*(imgs[i]+1.0)) # we ensure that there is no negative value wehn displaying images\n",
+    "\n",
+    "print(\"Norm L2(imgs[2]-imgs[0])\"+str(np.linalg.norm(imgs[2]-imgs[0])))\n",
+    "print(\"Norm L2(imgs[2]) \"+str(np.linalg.norm(imgs[2])))\n",
+    "print(\"Norm L2(count pixels) \"+str(math.sqrt(np.size(imgs[2][imgs[2]==-1]))))"
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "### Expe parameters\n",
     "\n",
@@ -332,56 +333,44 @@
     "For both binary generator and ternary generator this distance is 35.\n",
     "\n",
     "We will then compute this distance using a neural network."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "batch_size=64\n",
     "epochs=5\n",
     "steps_per_epoch=6400"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "generator = ternary_generator   #binary_generator, ternary_generator\n",
     "activation = FullSort #ReLU, MaxMin, GroupSort"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "### Build lipschitz Model"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 10,
-   "source": [
-    "K.clear_session()\n",
-    "## please note that the previous helper function has the same behavior as the following code:\n",
-    "inputs = Input((img_size, img_size, 1))\n",
-    "x = Flatten()(inputs)\n",
-    "x = SpectralDense(128, activation=FullSort())(x)\n",
-    "x = SpectralDense(64, activation=FullSort())(x)\n",
-    "x = SpectralDense(32, activation=FullSort())(x)\n",
-    "y = FrobeniusDense(1, activation=None)(x)\n",
-    "wass = Model(inputs=inputs, outputs=y)\n",
-    "wass.summary()"
-   ],
+   "metadata": {},
    "outputs": [
     {
-     "output_type": "stream",
      "name": "stderr",
+     "output_type": "stream",
      "text": [
       "2021-09-08 18:20:38.075170: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set\n",
       "2021-09-08 18:20:38.076265: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcuda.so.1\n",
@@ -430,8 +419,8 @@
      ]
     },
     {
-     "output_type": "stream",
      "name": "stdout",
+     "output_type": "stream",
      "text": [
       "Model: \"model\"\n",
       "_________________________________________________________________\n",
@@ -456,45 +445,52 @@
      ]
     }
    ],
-   "metadata": {}
+   "source": [
+    "K.clear_session()\n",
+    "## please note that the previous helper function has the same behavior as the following code:\n",
+    "inputs = Input((img_size, img_size, 1))\n",
+    "x = Flatten()(inputs)\n",
+    "x = SpectralDense(128, activation=FullSort())(x)\n",
+    "x = SpectralDense(64, activation=FullSort())(x)\n",
+    "x = SpectralDense(32, activation=FullSort())(x)\n",
+    "y = FrobeniusDense(1, activation=None)(x)\n",
+    "wass = Model(inputs=inputs, outputs=y)\n",
+    "wass.summary()"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "optimizer = Adam(lr=0.01)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 12,
-   "source": [
-    "wass.compile(loss=negative_KR, optimizer=optimizer, metrics=[KR])  "
-   ],
+   "metadata": {},
    "outputs": [],
-   "metadata": {}
+   "source": [
+    "wass.compile(loss=HKR(alpha=0), optimizer=optimizer, metrics=[KR])"
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "### Learn on toy dataset"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 13,
-   "source": [
-    "wass.fit_generator( generator(batch_size,(img_size,img_size,1),frac=frac_value),\n",
-    "                steps_per_epoch=steps_per_epoch// batch_size,\n",
-    "                epochs=epochs,verbose=1)"
-   ],
+   "metadata": {},
    "outputs": [
     {
-     "output_type": "stream",
      "name": "stderr",
+     "output_type": "stream",
      "text": [
       "/home/thibaut.boissin/envs/tf24/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py:1844: UserWarning: `Model.fit_generator` is deprecated and will be removed in a future version. Please use `Model.fit`, which supports generators.\n",
       "  warnings.warn('`Model.fit_generator` is deprecated and '\n",
@@ -503,8 +499,8 @@
      ]
     },
     {
-     "output_type": "stream",
      "name": "stdout",
+     "output_type": "stream",
      "text": [
       "Epoch 1/5\n",
       "100/100 [==============================] - 6s 50ms/step - loss: -24.9882 - KR: 24.9882\n",
@@ -519,30 +515,37 @@
      ]
     },
     {
-     "output_type": "execute_result",
      "data": {
       "text/plain": [
        "<tensorflow.python.keras.callbacks.History at 0x7f6a6b73cfd0>"
       ]
      },
+     "execution_count": 13,
      "metadata": {},
-     "execution_count": 13
+     "output_type": "execute_result"
     }
    ],
-   "metadata": {}
+   "source": [
+    "wass.fit_generator( generator(batch_size,(img_size,img_size,1),frac=frac_value),\n",
+    "                steps_per_epoch=steps_per_epoch// batch_size,\n",
+    "                epochs=epochs,verbose=1)"
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "As we can see the loss converge to the value 35 which is the wasserstein distance between the two distributions (square and non-square)."
-   ],
-   "metadata": {}
+   ]
   }
  ],
  "metadata": {
+  "interpreter": {
+   "hash": "e585d72a124540032141457729caea4129d351be49f1f69f41c00c4f8476abb5"
+  },
   "kernelspec": {
-   "name": "python3",
-   "display_name": "Python 3.7.11 64-bit ('tf24': venv)"
+   "display_name": "Python 3.7.11 64-bit ('tf24': venv)",
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -555,11 +558,8 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.7.11"
-  },
-  "interpreter": {
-   "hash": "e585d72a124540032141457729caea4129d351be49f1f69f41c00c4f8476abb5"
   }
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
\ No newline at end of file
+}
diff --git a/doc/source/conf.py b/doc/source/conf.py
index d54f3192..eeef0cd4 100644
--- a/doc/source/conf.py
+++ b/doc/source/conf.py
@@ -31,7 +31,7 @@
 )
 
 # The full version, including alpha/beta/rc tags
-release = "1.2.0"
+release = "1.3.0"
 
 # -- General configuration ---------------------------------------------------
 
@@ -41,9 +41,7 @@
 extensions = [
     "sphinx.ext.autodoc",
     "sphinx.ext.napoleon",
-    "sphinx_markdown_builder",
     "sphinx_rtd_theme",
-    "recommonmark",
 ]
 
 autoclass_content = "both"
@@ -84,18 +82,6 @@
     "titles_only": False,
 }
 
-html_sidebars = {
-    "**": ["globaltoc.html", "relations.html", "sourcelink.html", "searchbox.html"]
-}
-
-html_static_path = ["_static"]
-
-html_context = {
-    "css_files": [
-        "_static/theme_overrides.css",
-    ],  # override wide tables in RTD theme
-}
-
 autodoc_member_order = [
     "bysource",
 ]
diff --git a/doc/source/deel.lip.metrics.rst b/doc/source/deel.lip.metrics.rst
new file mode 100644
index 00000000..8f12b883
--- /dev/null
+++ b/doc/source/deel.lip.metrics.rst
@@ -0,0 +1,7 @@
+deel.lip.metrics module
+=======================
+
+.. automodule:: deel.lip.metrics
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/doc/source/deel.lip.rst b/doc/source/deel.lip.rst
index c218ce80..f7664493 100644
--- a/doc/source/deel.lip.rst
+++ b/doc/source/deel.lip.rst
@@ -13,8 +13,10 @@ Submodules
    deel.lip.initializers
    deel.lip.layers
    deel.lip.losses
+   deel.lip.metrics
    deel.lip.model
    deel.lip.normalizers
+   deel.lip.unconstrained_layers
    deel.lip.utils
 
 Module contents
diff --git a/doc/source/deel.lip.unconstrained_layers.rst b/doc/source/deel.lip.unconstrained_layers.rst
new file mode 100644
index 00000000..f32476c3
--- /dev/null
+++ b/doc/source/deel.lip.unconstrained_layers.rst
@@ -0,0 +1,7 @@
+deel.lip.unconstrained\_layers module
+=====================================
+
+.. automodule:: deel.lip.unconstrained_layers
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/setup.cfg b/setup.cfg
index e75e1925..88b3d26b 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -4,7 +4,8 @@ max-line-length = 88
 extend-ignore =
     # See https://github.com/PyCQA/pycodestyle/issues/373
     E203, E231
-
+per-file-ignores = 
+	*/__init__.py: F401
 
 [tox:tox]
 envlist = py{36,37,38},py{36,37,38}-lint
diff --git a/setup.py b/setup.py
index 5cfa5b73..88e729de 100644
--- a/setup.py
+++ b/setup.py
@@ -30,7 +30,7 @@
 
 setuptools.setup(
     name="deel-lip",
-    version="1.2.0",
+    version="1.3.0",
     author=", ".join(["Mathieu SERRURIER", "Franck MAMALET", "Thibaut BOISSIN"]),
     author_email=", ".join(
         [


Test patch:
diff --git a/tests/test_activations.py b/tests/test_activations.py
new file mode 100644
index 00000000..6411c007
--- /dev/null
+++ b/tests/test_activations.py
@@ -0,0 +1,111 @@
+from unittest import TestCase
+import tensorflow as tf
+from tensorflow.keras.models import Sequential, load_model
+from tensorflow.keras.layers import Input
+from tensorflow.keras.optimizers import SGD
+from tensorflow.keras.losses import CategoricalCrossentropy
+from deel.lip.activations import GroupSort
+import os
+import numpy as np
+
+
+def check_serialization(layer):
+    m = Sequential([Input(10), layer])
+    m.compile(optimizer=SGD(), loss=CategoricalCrossentropy(from_logits=False))
+    name = layer.__class__.__name__
+    path = os.path.join("logs", "losses", name)
+    x = tf.random.uniform((255, 10), -10, 10)
+    y1 = m(x)
+    m.save(path)
+    m2 = load_model(path, compile=True)
+    y2 = m2(x)
+    np.testing.assert_allclose(y1.numpy(), y2.numpy())
+
+
+class TestGroupSort(TestCase):
+    def test_simple(self):
+        gs = GroupSort(n=2)
+        check_serialization(gs)
+        gs = GroupSort(n=5)
+        check_serialization(gs)
+
+    def test_flat_input(self):
+        gs = GroupSort(n=2)
+        x = tf.convert_to_tensor(
+            [
+                [1, 2, 3, 4],
+                [4, 3, 2, 1],
+                [1, 2, 1, 2],
+                [2, 1, 2, 1],
+            ],
+            dtype=tf.float32,
+        )
+        gs.build(tf.TensorShape((None, 4)))
+        y = gs(x).numpy()
+        y_t = [
+            [1.0, 2.0, 3.0, 4.0],
+            [3.0, 4.0, 1.0, 2.0],
+            [1.0, 2.0, 1.0, 2.0],
+            [1.0, 2.0, 1.0, 2.0],
+        ]
+        np.testing.assert_equal(y, y_t)
+
+    def test_gs4(self):
+        gs = GroupSort(n=4)
+        x = tf.convert_to_tensor(
+            [
+                [1, 2, 3, 4],
+                [4, 3, 2, 1],
+                [1, 2, 1, 2],
+                [2, 1, 2, 1],
+            ],
+            dtype=tf.float32,
+        )
+        gs.build(tf.TensorShape((None, 4)))
+        y = gs(x).numpy()
+        y_t = [
+            [1.0, 2.0, 3.0, 4.0],
+            [1.0, 2.0, 3.0, 4.0],
+            [1.0, 1.0, 2.0, 2.0],
+            [1.0, 1.0, 2.0, 2.0],
+        ]
+        np.testing.assert_equal(y, y_t)
+
+    def test_img_input(self):
+        gs = GroupSort(n=2)
+        x = tf.convert_to_tensor(
+            [
+                [1, 2, 3, 4],
+                [4, 3, 2, 1],
+                [1, 2, 1, 2],
+                [2, 1, 2, 1],
+            ],
+            dtype=tf.float32,
+        )
+        x = tf.repeat(tf.expand_dims(tf.repeat(tf.expand_dims(x, 1), 28, 1), 1), 28, 1)
+        gs.build(tf.TensorShape((None, 28, 28, 4)))
+        y = gs(x).numpy()
+        y_t = [
+            [1.0, 2.0, 3.0, 4.0],
+            [3.0, 4.0, 1.0, 2.0],
+            [1.0, 2.0, 1.0, 2.0],
+            [1.0, 2.0, 1.0, 2.0],
+        ]
+        y_t = tf.repeat(
+            tf.expand_dims(tf.repeat(tf.expand_dims(y_t, 1), 28, 1), 1), 28, 1
+        )
+        np.testing.assert_equal(y, y_t)
+
+    def test_idempotence(self):
+        gs = GroupSort(n=2)
+        x = tf.random.uniform((255, 16), -10, 10)
+        gs.build(tf.TensorShape((None, 16)))
+        y1 = gs(x)
+        y2 = gs(y1)
+        np.testing.assert_equal(y1.numpy(), y2.numpy())
+        gs = GroupSort(n=2)
+        x = tf.random.uniform((255, 16), -10, 10)
+        gs.build(tf.TensorShape((None, 16)))
+        y1 = gs(x)
+        y2 = gs(y1)
+        np.testing.assert_equal(y1.numpy(), y2.numpy())
diff --git a/tests/test_initializers.py b/tests/test_initializers.py
index 99bc89d9..9e45044e 100644
--- a/tests/test_initializers.py
+++ b/tests/test_initializers.py
@@ -20,12 +20,16 @@
 class MyTestCase(unittest.TestCase):
     def test_bjorck_initializer(self):
         input_shape = (5,)
-        model = Sequential([Dense(4, kernel_initializer=SpectralInitializer(30, 15))])
-        self._test_model(model, input_shape)
-        model = Sequential([Dense(100, kernel_initializer=SpectralInitializer(30, 15))])
-        self._test_model(model, input_shape)
+        model = Sequential(
+            [Dense(4, kernel_initializer=SpectralInitializer(1e-6, 1e-6))]
+        )
+        self._test_model(model, input_shape, orthogonal_test=True)
+        model = Sequential(
+            [Dense(100, kernel_initializer=SpectralInitializer(1e-6, None))]
+        )
+        self._test_model(model, input_shape, orthogonal_test=False)
 
-    def _test_model(self, model, input_shape):
+    def _test_model(self, model, input_shape, orthogonal_test=True):
         batch_size = 1000
         # clear session to avoid side effects from previous train
         K.clear_session()
@@ -40,7 +44,10 @@ def _test_model(self, model, input_shape):
             full_matrices=False,
             compute_uv=False,
         ).numpy()
-        np.testing.assert_allclose(sigmas, np.ones_like(sigmas), 1e-6, 0)
+        if orthogonal_test:
+            np.testing.assert_allclose(sigmas, np.ones_like(sigmas), 1e-5, 0)
+        else:
+            np.testing.assert_almost_equal(sigmas.max(), 1.0, 5)
 
 
 if __name__ == "__main__":
diff --git a/tests/test_layers.py b/tests/test_layers.py
index 622cc0ee..2a32d56b 100644
--- a/tests/test_layers.py
+++ b/tests/test_layers.py
@@ -440,7 +440,7 @@ def test_spectral_dense(self):
             [
                 dict(
                     layer_type=SpectralDense,
-                    layer_params={"units": 3, "use_bias": False, "niter_spectral": 5},
+                    layer_params={"units": 3, "use_bias": False},
                     batch_size=250,
                     steps_per_epoch=125,
                     epochs=5,
@@ -451,7 +451,7 @@ def test_spectral_dense(self):
                 ),
                 dict(
                     layer_type=SpectralDense,
-                    layer_params={"units": 4, "niter_spectral": 5},
+                    layer_params={"units": 4},
                     batch_size=250,
                     steps_per_epoch=125,
                     epochs=5,
@@ -462,7 +462,7 @@ def test_spectral_dense(self):
                 ),
                 dict(
                     layer_type=SpectralDense,
-                    layer_params={"units": 4, "niter_spectral": 5},
+                    layer_params={"units": 4},
                     batch_size=250,
                     steps_per_epoch=125,
                     epochs=5,
diff --git a/tests/test_losses.py b/tests/test_losses.py
index f24410d0..a547dc97 100644
--- a/tests/test_losses.py
+++ b/tests/test_losses.py
@@ -7,7 +7,6 @@
 from tensorflow.keras.models import load_model
 from deel.lip.losses import (
     KR,
-    negative_KR,
     HingeMargin,
     HKR,
     MulticlassKR,
@@ -15,6 +14,7 @@
     MulticlassHKR,
     MultiMargin,
 )
+from deel.lip.utils import process_labels_for_multi_gpu
 import os
 import numpy as np
 
@@ -30,7 +30,7 @@ def get_gaussian_data(n=500, mean1=1.0, mean2=-1.0):
     x1 = tf.random.normal((n, 1), mean=mean1, stddev=0.1)
     x2 = tf.random.normal((n, 1), mean=mean2, stddev=0.1)
     y_pred = tf.concat([x1, x2], axis=0)
-    y_true = tf.concat([tf.ones(n), tf.zeros(n)], axis=0)
+    y_true = tf.concat([tf.ones((n, 1)), tf.zeros((n, 1))], axis=0)
     return y_pred, y_true
 
 
@@ -44,53 +44,33 @@ def check_serialization(nb_class, loss):
     m2(tf.random.uniform((255, 10)))
 
 
+def binary_tf_data(x):
+    """Return a TF float32 tensor of shape [N, 1] from a list/np.array of shape [N]"""
+    return tf.expand_dims(tf.constant(x, dtype=tf.float32), axis=-1)
+
+
 class Test(TestCase):
     def test_kr_loss(self):
-        loss = KR
-        y_pred, y_true = get_gaussian_data(20000)
+        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])
         loss_val = loss(y_true, y_pred).numpy()
-        np.testing.assert_approx_equal(
-            loss_val, 2.0, 1, "test failed when y_true has shape (bs, )"
-        )
-        loss_val_2 = loss(
-            tf.expand_dims(y_true, axis=-1), tf.expand_dims(y_pred, axis=-1)
-        ).numpy()
-        np.testing.assert_approx_equal(
-            loss_val_2, 2.0, 1, "test failed when y_true has shape (bs, 1)"
-        )
-        loss_val_3 = loss(tf.cast(y_true, dtype=tf.int32), y_pred).numpy()
-        np.testing.assert_approx_equal(
-            loss_val_3, 2.0, 1, "test failed when y_true has dtype int32"
-        )
-        y_true2 = tf.where(y_true == 1.0, 1.0, -1.0)
-        loss_val_4 = loss(y_true2, y_pred).numpy()
-        np.testing.assert_approx_equal(
-            loss_val_4, 2.0, 1, "test failed when labels are in (1, -1)"
-        )
-        check_serialization(1, loss)
+        self.assertEqual(loss_val, np.float32(1), "KR loss must be equal to 1")
 
-    def test_neg_kr_loss(self):
-        loss = negative_KR
         y_pred, y_true = get_gaussian_data(20000)
         loss_val = loss(y_true, y_pred).numpy()
         np.testing.assert_approx_equal(
-            loss_val, -2.0, 1, "test failed when y_true has shape (bs, )"
-        )
-        loss_val_2 = loss(
-            tf.expand_dims(y_true, axis=-1), tf.expand_dims(y_pred, axis=-1)
-        ).numpy()
-        np.testing.assert_approx_equal(
-            loss_val_2, -2.0, 1, "test failed when y_true has shape (bs, 1)"
+            loss_val, 2.0, 1, "test failed when y_true has shape (bs, 1)"
         )
         loss_val_3 = loss(tf.cast(y_true, dtype=tf.int32), y_pred).numpy()
-        np.testing.assert_approx_equal(
-            loss_val_3, -2.0, 1, "test failed when y_true has dtype int32"
+        self.assertEqual(
+            loss_val_3, loss_val, "test failed when y_true has dtype int32"
         )
         y_true2 = tf.where(y_true == 1.0, 1.0, -1.0)
         loss_val_4 = loss(y_true2, y_pred).numpy()
-        np.testing.assert_approx_equal(
-            loss_val_4, -2.0, 1, "test failed when labels are in (1, -1)"
-        )
+        self.assertEqual(loss_val_4, loss_val, "test failed when labels are in (1, -1)")
+
         check_serialization(1, loss)
 
     def test_hkr_loss(self):
@@ -103,53 +83,42 @@ def test_hinge_margin_loss(self):
         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_val = loss(y_true, y_pred).numpy()
-        np.testing.assert_almost_equal(
-            loss_val, 0.66, 1, "test failed when y_true has shape (bs, )"
-        )
-        loss_val_2 = loss(
-            tf.expand_dims(y_true, axis=-1), tf.expand_dims(y_pred, axis=-1)
-        ).numpy()
-        np.testing.assert_almost_equal(
-            loss_val_2, 0.66, 1, "test failed when y_true has shape (bs, 1)"
-        )
+        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()
-        np.testing.assert_almost_equal(
-            loss_val_3, 0.66, 1, "test failed when y_true has dtype int32"
+        self.assertEqual(
+            loss_val_3, np.float32(4 / 6), "Hinge loss must be equal to 4/6"
         )
         y_true2 = tf.where(y_true == 1.0, 1.0, -1.0)
         loss_val_4 = loss(y_true2, y_pred).numpy()
-        np.testing.assert_almost_equal(
-            loss_val_4, 0.66, 1, "test failed when labels are in (1, -1)"
+        self.assertEqual(
+            loss_val_4, np.float32(4 / 6), "Hinge loss must be equal to 4/6"
         )
         check_serialization(1, loss)
 
     def test_kr_multiclass_loss(self):
-        multiclass_kr = MulticlassKR()
-        kr = 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])
-        l_single = kr(y_true, y_pred).numpy()
-        l_multi = multiclass_kr(y_true, y_pred).numpy()
-        self.assertEqual(
-            l_single,
-            l_multi,
-            "KR multiclass must yield the same "
-            "results when given a single class "
-            "vector",
+        multiclass_kr = MulticlassKR(reduction="auto")
+        y_true = tf.one_hot([0, 0, 0, 1, 1, 2], 3)
+        y_pred = np.float32(
+            [
+                [2, 0.2, -0.5],
+                [-1, -1.2, 0.3],
+                [0.8, 2, 0],
+                [0, 1, -0.5],
+                [2.4, -0.4, -1.1],
+                [-0.1, -1.7, 0.6],
+            ]
         )
+        loss_val = multiclass_kr(y_true, y_pred).numpy()
+        np.testing.assert_allclose(loss_val, np.float32(761 / 1800))
+
         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 = multiclass_kr(y_true, y_pred).numpy()
         loss_val_2 = multiclass_kr(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"
-        )
-        y_true2 = tf.where(y_true == 1.0, 1.0, -1.0)
-        loss_val_3 = multiclass_kr(y_true2, y_pred).numpy()
-        np.testing.assert_almost_equal(
-            loss_val_3, loss_val, 1, "test failed when labels are in (1, -1)"
+        self.assertEqual(
+            loss_val_2, loss_val, "test failed when y_true has dtype int32"
         )
         check_serialization(1, multiclass_kr)
 
@@ -160,6 +129,7 @@ def test_hinge_multiclass_loss(self):
         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,
@@ -184,44 +154,42 @@ def test_hinge_multiclass_loss(self):
         y_pred = tf.random.normal((n_items, n_class))
         loss_val = multiclass_hinge(y_true, y_pred).numpy()
         loss_val_2 = multiclass_hinge(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"
+        self.assertEqual(
+            loss_val_2, loss_val, "test failed when y_true has dtype int32"
         )
         y_true2 = tf.where(y_true == 1.0, 1.0, -1.0)
         loss_val_3 = multiclass_hinge(y_true2, y_pred).numpy()
-        np.testing.assert_almost_equal(
-            loss_val_3, loss_val_2, 1, "test failed when labels are in (1, -1)"
+        self.assertEqual(
+            loss_val_3, loss_val_2, "test failed when labels are in (1, -1)"
         )
         check_serialization(1, multiclass_hinge)
 
     def test_hkr_multiclass_loss(self):
-        multiclass_hkr = MulticlassHKR(5, 1.0)
-        hkr_binary = HKR(5.0, 1.0)
-        # testing with an other value for eps ensure that eps has no influence
-        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 = hkr_binary(y_true, y_pred).numpy()
-        l_multi = multiclass_hkr(y_true, y_pred).numpy()
-        self.assertEqual(
-            l_single,
-            l_multi,
-            "hkr multiclass must yield the same "
-            "results when given a single class "
-            "vector",
+        multiclass_hkr = MulticlassHKR(5.0, 1.0)
+        y_true = tf.one_hot([0, 0, 0, 1, 1, 2], 3)
+        y_pred = np.float32(
+            [
+                [2, 0.2, -0.5],
+                [-1, -1.2, 0.3],
+                [0.8, 2, 0],
+                [0, 1, -0.5],
+                [2.4, -0.4, -1.1],
+                [-0.1, -1.7, 0.6],
+            ]
         )
+        loss_val = multiclass_hkr(y_true, y_pred).numpy()
+        np.testing.assert_allclose(loss_val, np.float32(1071 / 200))
+
         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 = multiclass_hkr(y_true, y_pred).numpy()
         loss_val_2 = multiclass_hkr(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"
-        )
-        y_true2 = tf.where(y_true == 1.0, 1.0, -1.0)
-        loss_val_3 = multiclass_hkr(y_true2, y_pred).numpy()
-        np.testing.assert_almost_equal(
-            loss_val_3, loss_val_2, 1, "test failed when labels are in (1, -1)"
+        self.assertEqual(
+            loss_val_2,
+            loss_val,
+            "test failed when y_true has dtype int32",
         )
         check_serialization(1, multiclass_hkr)
 
@@ -233,12 +201,301 @@ def test_multi_margin_loss(self):
         y_pred = tf.random.normal((n_items, n_class))
         loss_val = multimargin_loss(y_true, y_pred).numpy()
         loss_val_2 = multimargin_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"
+        self.assertEqual(
+            loss_val_2, loss_val, "test failed when y_true has dtype int32"
         )
         y_true2 = tf.where(y_true == 1.0, 1.0, -1.0)
         loss_val_3 = multimargin_loss(y_true2, y_pred).numpy()
-        np.testing.assert_almost_equal(
-            loss_val_3, loss_val_2, 1, "test failed when labels are in (1, -1)"
+        self.assertEqual(
+            loss_val_3, loss_val_2, "test failed when labels are in (1, -1)"
         )
         check_serialization(1, multimargin_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))
+
+        losses = (
+            KR(reduction="none"),
+            HingeMargin(0.7, reduction="none"),
+            HKR(alpha=2.5, reduction="none"),
+        )
+
+        expected_loss_values = (
+            np.array([1.0, 2.2, -0.2, 1.4, 2.6, 0.8, -0.4, 1.8]),
+            np.array([0.2, 0, 0.8, 0, 0, 0.3, 0.9, 0]),
+            np.array([0.25, -2.2, 2.95, -0.65, -2.6, 0.7, 3.4, -1.55]),
+        )
+
+        for loss, expected_loss_val in zip(losses, expected_loss_values):
+            loss_val = loss(y_true, y_pred)
+            np.testing.assert_allclose(
+                loss_val,
+                expected_loss_val,
+                rtol=5e-6,
+                atol=5e-6,
+                err_msg=f"Loss {loss.name} failed",
+            )
+
+    def test_no_reduction_multiclass_losses(self):
+        """
+        Assert multi-class losses without reduction. Four losses are tested on hardcoded
+        y_true/y_pred of shape [8 elements, 3 classes]: MulticlassKR, MulticlassHinge,
+        MultiMargin and MulticlassHKR.
+        """
+        y_true = tf.one_hot([0, 0, 0, 1, 1, 1, 2, 2], 3)
+        y_pred = np.array(
+            [
+                [1.5, 0.2, -0.5],
+                [1, -1.2, 0.3],
+                [0.8, 2, -2],
+                [0, 1, -0.5],
+                [-0.8, 2, 0],
+                [2.4, -0.4, -1.1],
+                [-0.1, -1.7, 0.6],
+                [1.2, 1.3, 2.5],
+            ],
+            dtype=np.float32,
+        )
+        losses = (
+            MulticlassKR(reduction="none"),
+            MulticlassHinge(reduction="none"),
+            MultiMargin(0.7, reduction="none"),
+            MulticlassHKR(alpha=2.5, min_margin=0.5, reduction="none"),
+        )
+
+        expected_loss_values = (
+            np.float64([326, 314, 120, 250, 496, -258, 396, 450]) / 225,
+            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,
+        )
+
+        for loss, expected_loss_val in zip(losses, expected_loss_values):
+            loss_val = loss(y_true, y_pred)
+            np.testing.assert_allclose(
+                loss_val,
+                expected_loss_val,
+                rtol=5e-6,
+                atol=5e-6,
+                err_msg=f"Loss {loss.name} failed",
+            )
+
+    def test_minibatches_binary_losses(self):
+        """
+        Assert binary losses with mini-batches, to simulate multi-GPU/TPU. Three losses
+        are tested using "sum" reduction: KR, HingeMargin and HKR.
+        Assertions are tested with:
+        - losses on hardcoded small data of size [8, 1] with full batch and mini-batches
+        - losses on large random data with full batch and mini-batches.
+        """
+        # Small hardcoded data of size [8, 1]
+        y_true = binary_tf_data([1, 1, 1, 1, 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])
+        y_true = process_labels_for_multi_gpu(y_true)
+
+        reduction = "sum"
+        losses = (
+            KR(multi_gpu=True, reduction=reduction),
+            HingeMargin(0.7, reduction=reduction),
+            HKR(alpha=2.5, multi_gpu=True, reduction=reduction),
+        )
+
+        expected_loss_values = (9.2, 2.2, 0.3)
+
+        # Losses are tested on full batch
+        for loss, expected_loss_val in zip(losses, expected_loss_values):
+            loss_val = loss(y_true, y_pred)
+            np.testing.assert_allclose(
+                loss_val,
+                expected_loss_val,
+                rtol=5e-6,
+                atol=5e-6,
+                err_msg=f"Loss {loss.name} failed for hardcoded full batch",
+            )
+
+        # Losses are tested on 3 mini-batches of size [3, 4, 1]
+        for loss, expected_loss_val in zip(losses, expected_loss_values):
+            loss_val = 0
+            loss_val += loss(y_true[:3], y_pred[:3])
+            loss_val += loss(y_true[3:7], y_pred[3:7])
+            loss_val += loss(y_true[7:], y_pred[7:])
+
+            np.testing.assert_allclose(
+                loss_val,
+                expected_loss_val,
+                rtol=5e-6,
+                atol=5e-6,
+                err_msg=f"Loss {loss.name} failed for hardcoded mini-batches",
+            )
+
+        # Large random data of size [10000, 1]
+        num_items = 10000
+        y_true = binary_tf_data(np.random.randint(2, size=num_items))
+        y_pred = tf.random.normal((num_items, 1))
+        y_true = process_labels_for_multi_gpu(y_true)
+
+        # Compare full batch loss and mini-batches loss
+        for loss in losses:
+            loss_val_full_batch = loss(y_true, y_pred)
+
+            i1 = num_items // 2
+            i2 = i1 + num_items // 4
+            i3 = i2 + num_items // 6
+            loss_val_minibatches = 0
+            loss_val_minibatches += loss(y_true[:i1], y_pred[:i1])
+            loss_val_minibatches += loss(y_true[i1:i2], y_pred[i1:i2])
+            loss_val_minibatches += loss(y_true[i2:i3], y_pred[i2:i3])
+            loss_val_minibatches += loss(y_true[i3:], y_pred[i3:])
+            np.testing.assert_allclose(
+                loss_val_full_batch,
+                loss_val_minibatches,
+                rtol=5e-6,
+                atol=5e-6,
+                err_msg=f"Loss {loss.name} failed for random mini-batches",
+            )
+
+    def test_minibatches_multiclass_losses(self):
+        """
+        Assert multiclass losses with mini-batches, to simulate multi-GPU/TPU. Four
+        losses are tested using "sum" reduction: MulticlassKR, MulticlassHinge,
+        MultiMargin and MulticlassHKR.
+        Assertions are tested with:
+        - losses on hardcoded small data of size [8, 3] with full batch and mini-batches
+        - losses on large random data with full batch and mini-batches.
+        """
+        # Small hardcoded data of size [8, 3]
+        y_true = tf.one_hot([0, 0, 0, 1, 1, 1, 2, 2], 3)
+        y_true = process_labels_for_multi_gpu(y_true)
+        y_pred = tf.constant(
+            [
+                [1.5, 0.2, -0.5],
+                [1, -1.2, 0.3],
+                [0.8, 2, -2],
+                [0, 1, -0.5],
+                [-0.8, 2, 0],
+                [2.4, -0.4, -1.1],
+                [-0.1, -1.7, 0.6],
+                [1.2, 1.3, 2.5],
+            ],
+            dtype=tf.float32,
+        )
+
+        # Losses are tested on full batch
+        reduction = "sum"
+        losses = (
+            MulticlassKR(multi_gpu=True, reduction=reduction),
+            MulticlassHinge(reduction=reduction),
+            MultiMargin(0.7, reduction=reduction),
+            MulticlassHKR(
+                alpha=2.5, min_margin=0.5, multi_gpu=True, reduction=reduction
+            ),
+        )
+
+        expected_loss_values = (
+            np.float32(698 / 75),
+            np.float32(43 / 6),
+            np.float32(9 / 5),
+            np.float32(152 / 75),
+        )
+
+        # Losses are tested on 3 mini-batches of size [3, 4, 1]
+        for loss, expected_loss_val in zip(losses, expected_loss_values):
+            loss_val = loss(y_true, y_pred)
+            np.testing.assert_allclose(
+                loss_val,
+                expected_loss_val,
+                rtol=5e-6,
+                atol=5e-6,
+                err_msg=f"Loss {loss.name} failed for hardcoded full batch",
+            )
+
+        # The dataset is now split into 3 mini-batches of size [3, 4, 1]
+        for loss, expected_loss_val in zip(losses, expected_loss_values):
+            loss_val = 0
+            loss_val += loss(y_true[:3], y_pred[:3])
+            loss_val += loss(y_true[3:7], y_pred[3:7])
+            loss_val += loss(y_true[7:], y_pred[7:])
+
+            np.testing.assert_allclose(
+                loss_val,
+                expected_loss_val,
+                rtol=5e-6,
+                atol=5e-6,
+                err_msg=f"Loss {loss.name} failed for hardcoded mini-batches",
+            )
+
+        # Large random data of size [10000, 1]
+        num_classes = 10
+        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))
+
+        # Compare full batch loss and mini-batches loss
+        for loss in losses:
+            loss_val_full_batch = loss(y_true, y_pred)
+
+            i1 = num_items // 2
+            i2 = i1 + num_items // 4
+            i3 = i2 + num_items // 6
+            loss_val_minibatches = 0
+            loss_val_minibatches += loss(y_true[:i1], y_pred[:i1])
+            loss_val_minibatches += loss(y_true[i1:i2], y_pred[i1:i2])
+            loss_val_minibatches += loss(y_true[i2:i3], y_pred[i2:i3])
+            loss_val_minibatches += loss(y_true[i3:], y_pred[i3:])
+            np.testing.assert_allclose(
+                loss_val_full_batch,
+                loss_val_minibatches,
+                rtol=5e-6,
+                atol=5e-6,
+                err_msg=f"Loss {loss.name} failed for random mini-batches",
+            )
+
+    def test_multilabel_losses(self):
+        """
+        Assert binary losses with multilabels.
+        Three losses are tested (KR, HingeMargin and HKR). We compare losses with three
+        separate binary classification and the corresponding multilabel problem.
+        """
+        # Create predictions and labels for 3 binary problems and the concatenated
+        # multilabel one.
+        y_pred1, y_true1 = get_gaussian_data(1000)
+        y_pred2, y_true2 = get_gaussian_data(1000)
+        y_pred3, y_true3 = get_gaussian_data(1000)
+        y_pred = tf.concat([y_pred1, y_pred2, y_pred3], axis=-1)
+        y_true = tf.concat([y_true1, y_true2, y_true3], axis=-1)
+
+        # Tested losses (different reductions, multi_gpu)
+        losses = (
+            KR(reduction="none", name="KR none"),
+            HingeMargin(0.4, reduction="none", name="hinge none"),
+            HKR(alpha=5.2, reduction="none", name="HKR none"),
+            KR(reduction="auto", name="KR auto"),
+            HingeMargin(0.6, reduction="auto", name="hinge auto"),
+            HKR(alpha=10, reduction="auto", name="HKR auto"),
+            KR(multi_gpu=True, reduction="sum", name="KR multi_gpu"),
+            HKR(alpha=3.2, multi_gpu=True, reduction="sum", name="HKR multi_gpu"),
+        )
+
+        # Compute loss values and assert that the multilabel value is equal to the mean
+        # of the three separate binary problems.
+        for loss in losses:
+            loss_val1 = loss(y_true1, y_pred1).numpy()
+            loss_val2 = loss(y_true2, y_pred2).numpy()
+            loss_val3 = loss(y_true3, y_pred3).numpy()
+            mean_loss_vals = (loss_val1 + loss_val2 + loss_val3) / 3
+
+            loss_val_multilabel = loss(y_true, y_pred).numpy()
+
+            np.testing.assert_allclose(
+                loss_val_multilabel,
+                mean_loss_vals,
+                rtol=5e-6,
+                atol=5e-6,
+                err_msg=f"Loss {loss.name} failed",
+            )
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
new file mode 100644
index 00000000..0e0cda72
--- /dev/null
+++ b/tests/test_metrics.py
@@ -0,0 +1,278 @@
+import unittest
+from collections import defaultdict
+from unittest import TestCase
+import tensorflow as tf
+from tensorflow.keras.models import Sequential, load_model
+from tensorflow.keras.layers import Input, Dense
+from tensorflow.keras.optimizers import SGD
+from deel.lip.metrics import (
+    CategoricalProvableRobustAccuracy,
+    BinaryProvableRobustAccuracy,
+    CategoricalProvableAvgRobustness,
+    BinaryProvableAvgRobustness,
+)
+import numpy as np
+import os
+
+
+def check_serialization(nb_class, loss, nb_classes):
+    n = 255
+    x = tf.random.uniform((n, 42))
+    y = tf.one_hot(
+        tf.convert_to_tensor(np.random.randint(nb_classes, size=n)), nb_classes
+    )
+    m = Sequential([Input(42), Dense(nb_class)])
+    m.compile(optimizer=SGD(), loss=loss)
+    l1 = m.evaluate(x, y)
+    name = loss.__class__.__name__
+    path = os.path.join("logs", "losses", name)
+    m.save(path)
+    m2 = load_model(path, compile=True)
+    l2 = m2.evaluate(x, y)
+    return l1, l2
+
+
+class Test(TestCase):
+    def test_serialization(self):
+        pra = CategoricalProvableRobustAccuracy(1, 1.0, disjoint_neurons=False)
+        l1, l2 = check_serialization(10, pra, 10)
+        self.assertEqual(
+            l1,
+            l2,
+            "serialization changed loss value for CategoricalProvableRobustAccuracy",
+        )
+        pra = BinaryProvableRobustAccuracy(1, 1.0)
+        l1, l2 = check_serialization(10, pra, 10)
+        self.assertEqual(
+            l1, l2, "serialization changed loss value for BinaryProvableRobustAccuracy"
+        )
+        par = CategoricalProvableAvgRobustness(
+            2.0, disjoint_neurons=False, negative_robustness=True
+        )
+        l1, l2 = check_serialization(10, par, 10)
+        self.assertEqual(
+            l1,
+            l2,
+            "serialization changed loss value for CategoricalProvableAvgRobustness",
+        )
+        par = BinaryProvableAvgRobustness(2.0, negative_robustness=True)
+        l1, l2 = check_serialization(10, par, 10)
+        self.assertEqual(
+            l1, l2, "serialization changed loss value for BinaryProvableAvgRobustness"
+        )
+
+    def test_provable_vs_adjusted(self):
+        n = 500
+        x = tf.random.normal((n, 10), mean=0.0, stddev=0.1)
+        y = tf.one_hot(tf.convert_to_tensor(np.random.randint(10, size=n)), 10)
+        for lip_cst in [1, 0.5, 2.0]:
+            for disjoint in [True, False]:
+                pr = CategoricalProvableAvgRobustness(
+                    lip_cst, disjoint, negative_robustness=False
+                )
+                ar = CategoricalProvableAvgRobustness(
+                    lip_cst, disjoint, negative_robustness=True
+                )
+                l1 = pr(y, y).numpy()
+                l2 = ar(y, y).numpy()
+                self.assertAlmostEqual(
+                    l1,
+                    l2,
+                    4,
+                    msg="provable and adjusted "
+                    "robustness must give same "
+                    "values when y_true==y_pred",
+                )
+                l1 = pr(y, x).numpy()
+                l2 = ar(y, x).numpy()
+                self.assertNotAlmostEqual(
+                    l1,
+                    l2,
+                    4,
+                    msg="provable and adjusted robustness must give different values"
+                    " when y_true!=y_pred",
+                )
+
+        x = tf.random.normal((n, 10), mean=0.0, stddev=0.1)
+        y = tf.one_hot(tf.convert_to_tensor(np.random.randint(10, size=n)), 10)
+        for lip_cst in [1.0, 0.5, 2.0]:
+            bpr = BinaryProvableAvgRobustness(lip_cst, negative_robustness=False)
+            bar = BinaryProvableAvgRobustness(lip_cst, negative_robustness=True)
+            l1 = bpr(y, y).numpy()
+            l2 = bar(y, y).numpy()
+            self.assertAlmostEqual(
+                l1,
+                l2,
+                4,
+                msg="provable and adjusted robustness must give same values when y_"
+                "true==y_pred",
+            )
+            l1 = pr(y, x).numpy()
+            l2 = ar(y, x).numpy()
+            self.assertNotAlmostEqual(
+                l1,
+                l2,
+                4,
+                msg="provable and adjusted robustness must give different values when "
+                "y_true!=y_pred",
+            )
+
+    def test_data_format(self):
+        pr = CategoricalProvableAvgRobustness(1.0, True, negative_robustness=False)
+        ar = CategoricalProvableAvgRobustness(1.0, True, negative_robustness=True)
+        n = 500
+        # check in multiclass
+        for metric in [pr, ar]:
+            y_pred = tf.random.normal((n, 10), mean=0.0, stddev=0.1)
+            y_true = tf.one_hot(tf.convert_to_tensor(np.random.randint(10, size=n)), 10)
+            metrics_values = []
+            for neg_val in [0.0, -1.0]:
+                y_pred_case = y_pred
+                y_true_case = tf.where(y_true > 0, 1.0, neg_val)
+                metrics_values.append(metric(y_true_case, y_pred_case).numpy())
+            self.assertTrue(
+                all([m == metrics_values[0] for m in metrics_values]),
+                "changing the data format must not change the metric value",
+            )
+        # check in binary
+        bpr = BinaryProvableAvgRobustness(1.0, negative_robustness=False)
+        bar = BinaryProvableAvgRobustness(1.0, negative_robustness=True)
+        for metric in [bpr, bar]:
+            y_pred = tf.random.normal((n,), mean=0.0, stddev=0.1)
+            y_true = tf.random.normal((n,), mean=0.0, stddev=0.1)
+            metrics_values = []
+            for neg_val in [0.0, -1.0]:
+                y_pred_case = y_pred
+                y_true_case = tf.where(y_true > 0, 1.0, neg_val)
+                metrics_values.append(metric(y_true_case, y_pred_case).numpy())
+            self.assertTrue(
+                all([m == metrics_values[0] for m in metrics_values]),
+                "changing the data format must not change the metric value",
+            )
+
+    def test_disjoint_neurons(self):
+        pr = CategoricalProvableAvgRobustness(1.0, False, negative_robustness=False)
+        ar = CategoricalProvableAvgRobustness(1.0, False, negative_robustness=True)
+        pdr = CategoricalProvableAvgRobustness(1.0, True, negative_robustness=False)
+        adr = CategoricalProvableAvgRobustness(1.0, True, negative_robustness=True)
+        n = 500
+        # check in multiclass
+        metrics_values = defaultdict(list)
+        y_pred = tf.random.normal((n, 10), mean=0.0, stddev=1.0)
+        y_true = tf.one_hot(tf.convert_to_tensor(np.random.randint(10, size=n)), 10)
+        for neg_val in [0.0, -1.0]:
+            y_pred_case = y_pred
+            y_true_case = tf.where(y_true > 0, 1.0, neg_val)
+            metrics_values["standard"].append(pdr(y_true_case, y_pred_case).numpy())
+            metrics_values["standard"].append(
+                pr(y_true_case, y_pred_case * (np.sqrt(2) / 2.0)).numpy()
+            )
+            metrics_values["negative_rob"].append(adr(y_true_case, y_pred_case).numpy())
+            metrics_values["negative_rob"].append(
+                ar(y_true_case, y_pred_case * (np.sqrt(2) / 2.0)).numpy()
+            )
+        print(metrics_values)
+        for metric_type in ["standard", "negative_rob"]:
+            for m in metrics_values[metric_type]:
+                self.assertAlmostEqual(
+                    m,
+                    metrics_values[metric_type][0],
+                    4,
+                    "the loss does not account disjoint_neuron properly in multiclass",
+                )
+        # check in binary
+        pr = BinaryProvableAvgRobustness(1.0, negative_robustness=False)
+        ar = BinaryProvableAvgRobustness(1.0, negative_robustness=True)
+        metrics_values = defaultdict(list)
+        y_pred = tf.random.normal((n,), mean=0.0, stddev=0.1)
+        y_true = tf.random.normal((n,), mean=0.0, stddev=0.1)
+        for neg_val in [0.0, -1.0]:
+            y_pred_case = y_pred
+            y_true_case = tf.where(y_true > 0, 1.0, neg_val)
+            # no sqrt(2)/2 here as the corrective factor works only in multiclass
+            metrics_values["standard"].append(pr(y_true_case, y_pred_case).numpy())
+            metrics_values["negative_rob"].append(ar(y_true_case, y_pred_case).numpy())
+        self.assertTrue(
+            all(
+                [m == metrics_values["standard"][0] for m in metrics_values["standard"]]
+            ),
+            "the loss does not account disjoint_neuron properly in binary",
+        )
+        self.assertTrue(
+            all(
+                [
+                    m == metrics_values["negative_rob"][0]
+                    for m in metrics_values["negative_rob"]
+                ]
+            ),
+            "the loss does not account disjoint_neuron properly in binary",
+        )
+
+    def test_hardcoded_values(self):
+        pr = CategoricalProvableRobustAccuracy(0.25, 1.0, False)
+        ar = CategoricalProvableAvgRobustness(1.0, False)
+        y_pred = tf.convert_to_tensor(
+            [
+                [1.0, 0.0, 0.0],  # good class & over 0.25
+                [0.1, 0.0, 0.0],  # good class & below 0.25
+                [0.0, 0.0, 1.0],  # wrong class & over 0.25
+                [0.0, 0.0, 0.1],  # wrong class & below 0.25
+            ]
+        )
+        y_true = tf.convert_to_tensor(
+            [
+                [1.0, 0.0, 0.0],
+                [1.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0],
+                [0.0, 1.0, 0.0],
+            ]
+        )
+        self.assertEqual(
+            pr(y_true, y_pred).numpy(),
+            0.25,
+            "ProvableRobustAccuracy did not "
+            "get the expected result with "
+            "hardcoded values ",
+        )
+        self.assertAlmostEqual(
+            ar(y_true, y_pred).numpy(),
+            0.25 * 1.1 / np.sqrt(2),
+            5,
+            "ProvableAvgRobustness did not "
+            "get the expected result with "
+            "hardcoded values ",
+        )
+        bpr = BinaryProvableRobustAccuracy(0.25, 1.0)
+        bar = BinaryProvableAvgRobustness(1.0, negative_robustness=False)
+        y_pred = tf.convert_to_tensor(
+            [
+                1.0,  # good class & over 0.25
+                0.1,  # good class & below 0.25
+                1.0,  # wrong class & over 0.25
+                0.1,  # wrong class & below 0.25
+                -1.0,
+                -0.1,
+                -1.0,
+                -0.1,
+            ]
+        )
+        y_true = tf.convert_to_tensor([1.0, 1.0, -1.0, -1.0, -1.0, -1.0, 1.0, 1.0])
+        self.assertEqual(
+            bpr(y_true, y_pred).numpy(),
+            0.25,
+            "BinaryProvableRobustAccuracy did not "
+            "get the expected result with "
+            "hardcoded values ",
+        )
+        self.assertAlmostEqual(
+            bar(y_true, y_pred).numpy(),
+            0.125 * (1.1 * 2),
+            5,
+            "BinaryProvableAvgRobustness did not "
+            "get the expected result with "
+            "hardcoded values ",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_models.py b/tests/test_models.py
new file mode 100644
index 00000000..4dcecf18
--- /dev/null
+++ b/tests/test_models.py
@@ -0,0 +1,112 @@
+"""These tests assert that deel.lip Sequential and Model objects behave as expected."""
+
+import warnings
+from unittest import TestCase
+import numpy as np
+import tensorflow as tf
+import tensorflow.keras.layers as kl
+from deel.lip import Sequential, Model, vanillaModel
+from deel.lip.layers import SpectralConv2D, SpectralDense, ScaledL2NormPooling2D
+from deel.lip.activations import GroupSort2
+
+
+def sequential_layers():
+    """Return list of layers for a Sequential model"""
+    return [
+        SpectralConv2D(6, 3, input_shape=(20, 20, 3)),
+        ScaledL2NormPooling2D(),
+        GroupSort2(),
+        kl.Flatten(),
+        SpectralDense(10),
+    ]
+
+
+def functional_input_output_tensors():
+    """Return input and output tensor of a Functional (hard-coded) model"""
+    inputs = tf.keras.Input((8, 8, 3))
+    x = SpectralConv2D(2, (3, 3), k_coef_lip=2.0)(inputs)
+    x = GroupSort2()(x)
+    x = ScaledL2NormPooling2D((2, 2), k_coef_lip=2.0)(x)
+    x1 = SpectralConv2D(2, (3, 3), k_coef_lip=2.0)(x)
+    x1 = GroupSort2()(x1)
+    x = kl.Add()([x, x1])
+    x = kl.Flatten()(x)
+    x = kl.Dense(4)(x)
+    x = SpectralDense(4, k_coef_lip=2.0)(x)
+    outputs = SpectralDense(2, k_coef_lip=2.0)(x)
+    return inputs, outputs
+
+
+class Test(TestCase):
+    def assert_model_outputs(self, model1, model2):
+        """Assert outputs are identical for both models on random inputs"""
+        x = np.random.random((10,) + model1.input_shape[1:]).astype(np.float32)
+        y1 = model1.predict(x)
+        y2 = model2.predict(x)
+        np.testing.assert_array_equal(y2, y1)
+
+    def test_keras_Sequential(self):
+        """Assert vanilla conversion of a tf.keras.Sequential model"""
+        model = tf.keras.Sequential(sequential_layers())
+        vanilla_model = vanillaModel(model)
+        self.assert_model_outputs(model, vanilla_model)
+
+    def test_deel_lip_Sequential(self):
+        """Assert vanilla conversion of a deel.lip.Sequential model"""
+        model = Sequential(sequential_layers())
+        vanilla_model = model.vanilla_export()
+        self.assert_model_outputs(model, vanilla_model)
+
+    def test_keras_Model(self):
+        """Assert vanilla conversion of a tf.keras.Model model"""
+        inputs, outputs = functional_input_output_tensors()
+        model = tf.keras.Model(inputs, outputs)
+        vanilla_model = vanillaModel(model)
+        self.assert_model_outputs(model, vanilla_model)
+
+    def test_deel_lip_Model(self):
+        """Assert vanilla conversion of a deel.lip.Model model"""
+        inputs, outputs = functional_input_output_tensors()
+        model = Model(inputs, outputs)
+        vanilla_model = model.vanilla_export()
+        self.assert_model_outputs(model, vanilla_model)
+
+    def test_warning_unsupported_1Lip_layers(self):
+        """Assert that some unsupported layers return a warning message that they are
+        not 1-Lipschitz and other supported layers don't raise a warning.
+        """
+
+        # Check that supported 1-Lipschitz layers do not raise a warning
+        supported_layers = [
+            kl.Input((32, 32, 3)),
+            kl.ReLU(),
+            kl.Activation("relu"),
+            kl.Softmax(),
+            kl.Flatten(),
+            kl.Reshape((10,)),
+            kl.MaxPool2D(),
+            SpectralDense(3),
+            ScaledL2NormPooling2D(),
+        ]
+        for lay in supported_layers:
+            with warnings.catch_warnings(record=True) as w:
+                Sequential([lay])
+                self.assertEqual(len(w), 0, f"Layer {lay.name} shouldn't raise warning")
+
+        # Check that unsupported layers raise a warning
+        unsupported_layers = [
+            kl.MaxPool2D(pool_size=3, strides=2),
+            kl.Add(),
+            kl.Concatenate(),
+            kl.Activation("gelu"),
+            kl.Dense(5),
+            kl.Conv2D(10, 3),
+            kl.UpSampling2D(),
+        ]
+        for lay in unsupported_layers:
+            with self.assertWarnsRegex(
+                Warning,
+                expected_regex="layer which is not a 1-Lipschitz",
+                msg=f"Layer {lay.name} should raise a warning.",
+            ):
+                Sequential([lay])
diff --git a/tests/test_normalizers.py b/tests/test_normalizers.py
new file mode 100644
index 00000000..3cd273ad
--- /dev/null
+++ b/tests/test_normalizers.py
@@ -0,0 +1,153 @@
+# 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 unittest
+
+import numpy as np
+import tensorflow as tf
+from deel.lip.normalizers import (
+    bjorck_normalization,
+    reshaped_kernel_orthogonalization,
+    spectral_normalization,
+)
+
+np.random.seed(42)
+
+
+class TestSpectralNorm(unittest.TestCase):
+    """Test of spectral normalization (power iteration) on Dense and Conv2D kernels."""
+
+    def test_spectral_normalization(self):
+        # Dense kernel
+        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)
+        kernel = np.random.normal(size=kernel_shape).astype("float32")
+        self._test_kernel(kernel)
+
+    def _test_kernel(self, kernel):
+        """Compare max singular value using power iteration and tf.linalg.svd"""
+        sigmas_svd = tf.linalg.svd(
+            np.reshape(kernel, (np.prod(kernel.shape[:-1]), kernel.shape[-1])),
+            full_matrices=False,
+            compute_uv=False,
+        ).numpy()
+        SVmax = np.max(sigmas_svd)
+
+        W_bar, _u, sigma = spectral_normalization(kernel, u=None, eps=1e-6)
+        # Test sigma is close to the one computed with svd first run @ 1e-1
+        np.testing.assert_approx_equal(
+            sigma, SVmax, 1, "test failed with kernel_shape " + str(kernel.shape)
+        )
+
+        W_bar, _u, sigma = spectral_normalization(kernel, u=_u, eps=1e-6)
+        W_bar, _u, sigma = spectral_normalization(kernel, u=_u, eps=1e-6)
+        # Test W_bar is reshaped correctly
+        np.testing.assert_equal(
+            W_bar.shape, (np.prod(kernel.shape[:-1]), kernel.shape[-1])
+        )
+        # Test sigma is close to the one computed with svd second run @ 1e-5
+        np.testing.assert_approx_equal(
+            sigma, SVmax, 3, "test failed with kernel_shape " + str(kernel.shape)
+        )
+        # Test if kernel is normalized by sigma
+        np.testing.assert_allclose(
+            np.reshape(W_bar, kernel.shape), kernel / sigmas_svd[0], 1e-2, 0
+        )
+
+
+class TestBjorckNormalization(unittest.TestCase):
+    """Test of Björck orthogonalization on Dense and Conv2D kernels."""
+
+    def test_bjorck_normalization(self):
+        # Dense kernel
+        kernel_shape = (15, 32)
+        kernel = np.random.normal(size=kernel_shape).astype("float32")
+        self._test_kernel(kernel)
+
+        # Conv kernel
+        kernel_shape = (64, 32)
+        kernel = np.random.normal(size=kernel_shape).astype("float32")
+        self._test_kernel(kernel)
+
+    def _test_kernel(self, kernel):
+        """Compare max singular value using power iteration and tf.linalg.svd"""
+        sigmas_svd = tf.linalg.svd(
+            np.reshape(kernel, (np.prod(kernel.shape[:-1]), kernel.shape[-1])),
+            full_matrices=False,
+            compute_uv=False,
+        ).numpy()
+        SVmax = np.max(sigmas_svd)
+
+        wbar = bjorck_normalization(kernel / SVmax, eps=1e-5)
+        sigmas_wbar_svd = tf.linalg.svd(
+            np.reshape(wbar, (np.prod(wbar.shape[:-1]), wbar.shape[-1])),
+            full_matrices=False,
+            compute_uv=False,
+        ).numpy()
+        # Test sigma is close to the one computed with svd first run @ 1e-1
+        np.testing.assert_allclose(
+            sigmas_wbar_svd, np.ones(sigmas_wbar_svd.shape), 1e-2, 0
+        )
+        # Test W_bar is reshaped correctly
+        np.testing.assert_equal(
+            wbar.shape, (np.prod(kernel.shape[:-1]), kernel.shape[-1])
+        )
+
+        # Test sigma is close to the one computed with svd second run @1e-5
+        wbar = bjorck_normalization(wbar, eps=1e-5)
+        sigmas_wbar_svd = tf.linalg.svd(
+            np.reshape(wbar, (np.prod(wbar.shape[:-1]), wbar.shape[-1])),
+            full_matrices=False,
+            compute_uv=False,
+        ).numpy()
+        np.testing.assert_allclose(
+            sigmas_wbar_svd, np.ones(sigmas_wbar_svd.shape), 1e-4, 0
+        )
+
+
+class TestRKO(unittest.TestCase):
+    """Test of RKO algorithm on Dense and Conv2D kernels."""
+
+    def test_reshaped_kernel_orthogonalization(self):
+        # Dense kernel
+        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)
+        kernel = np.random.normal(size=kernel_shape).astype("float32")
+        self._test_kernel(kernel)
+
+    def _test_kernel(self, kernel):
+        """Compare max singular value using power iteration and tf.linalg.svd"""
+        sigmas_svd = tf.linalg.svd(
+            np.reshape(kernel, (np.prod(kernel.shape[:-1]), kernel.shape[-1])),
+            full_matrices=False,
+            compute_uv=False,
+        ).numpy()
+        SVmax = np.max(sigmas_svd)
+
+        W_bar, _, sigma = reshaped_kernel_orthogonalization(
+            kernel, u=None, adjustment_coef=1.0, eps_spectral=1e-5, eps_bjorck=1e-5
+        )
+        # Test W_bar is reshaped correctly
+        np.testing.assert_equal(W_bar.shape, kernel.shape)
+        # Test RKO sigma is close to max(svd)
+        np.testing.assert_approx_equal(
+            sigma, SVmax, 1, "test failed with kernel_shape " + str(kernel.shape)
+        )
+        sigmas_wbar_svd = tf.linalg.svd(
+            np.reshape(W_bar, (np.prod(W_bar.shape[:-1]), W_bar.shape[-1])),
+            full_matrices=False,
+            compute_uv=False,
+        ).numpy()
+        # Test if SVs of W_bar are close to one
+        np.testing.assert_allclose(
+            sigmas_wbar_svd, np.ones(sigmas_wbar_svd.shape), 1e-2, 0
+        )
diff --git a/tests/test_unconstrained_layers.py b/tests/test_unconstrained_layers.py
new file mode 100644
index 00000000..9a9dc6d9
--- /dev/null
+++ b/tests/test_unconstrained_layers.py
@@ -0,0 +1,192 @@
+import unittest
+
+import numpy as np
+import tensorflow as tf
+from deel.lip.unconstrained_layers import PadConv2D
+from deel.lip.model import vanillaModel
+from deel.lip.utils import _padding_circular
+
+from .test_layers import generate_k_lip_model
+
+
+class TestPadConv2D(unittest.TestCase):
+    def test_PadConv2D(self):
+        """Main test for PadConv2D: tests on padding, predict and vanilla export."""
+
+        all_paddings = ("circular", "constant", "symmetric", "reflect", "same", "valid")
+        layer_params = {
+            "filters": 2,
+            "kernel_size": (3, 3),
+            "use_bias": False,
+            "padding": "valid",
+        }
+        batch_size = 250
+        kernel_input_shapes = [
+            [(3, 3), 2, (5, 5, 1)],
+            [(3, 3), 2, (5, 5, 5)],
+            [(7, 7), 2, (7, 7, 2)],
+            [3, 2, (5, 5, 128)],
+        ]
+
+        for k_shape, filters, input_shape in kernel_input_shapes:
+            layer_params["kernel_size"] = k_shape
+            layer_params["filters"] = filters
+            for pad in ["circular", "constant", "symmetric", "reflect"]:
+                self._test_padding(pad, input_shape, batch_size, k_shape)
+            for pad in all_paddings:
+                self._test_predict(layer_params, pad, input_shape, batch_size)
+            for pad in all_paddings:
+                self._test_vanilla(layer_params, pad, input_shape, batch_size)
+
+    def pad_input(self, x, padding, kernel_size):
+        """Pad an input tensor x with corresponding padding, based on kernel size."""
+        if isinstance(kernel_size, (int, float)):
+            kernel_size = [kernel_size, kernel_size]
+        if padding.lower() in ["same", "valid"]:
+            return x
+        elif padding.lower() in ["constant", "reflect", "symmetric"]:
+            p_vert, p_hor = kernel_size[0] // 2, kernel_size[1] // 2
+            pad_sizes = [[0, 0], [p_vert, p_vert], [p_hor, p_hor], [0, 0]]
+            return tf.pad(x, pad_sizes, padding)
+        elif padding.lower() in ["circular"]:
+            p_vert, p_hor = kernel_size[0] // 2, kernel_size[1] // 2
+            return _padding_circular(x, (p_vert, p_hor))
+
+    def compare(self, x, x_ref, index_x=[], index_x_ref=[]):
+        """Compare a tensor and its padded version, based on index_x and ref."""
+        x_cropped = x[:, index_x[0] : index_x[1], index_x[3] : index_x[4], :][
+            :, :: index_x[2], :: index_x[5], :
+        ]
+        if index_x_ref[0] is None:  # compare with 0
+            np.testing.assert_allclose(x_cropped, np.zeros(x_cropped.shape), 1e-2, 0)
+        else:
+            np.testing.assert_allclose(
+                x_cropped,
+                x_ref[
+                    :,
+                    index_x_ref[0] : index_x_ref[1],
+                    index_x_ref[3] : index_x_ref[4],
+                    :,
+                ][:, :: index_x_ref[2], :: index_x_ref[5], :],
+                1e-2,
+                0,
+            )
+
+    def _test_padding(self, padding_tested, input_shape, batch_size, kernel_size):
+        """Test different padding types: assert values in original and padded tensors"""
+        kernel_size_list = kernel_size
+        if isinstance(kernel_size, (int, float)):
+            kernel_size_list = [kernel_size, kernel_size]
+
+        x = np.random.normal(size=(batch_size,) + input_shape).astype("float32")
+        x_pad = self.pad_input(x, padding_tested, kernel_size)
+        p_vert, p_hor = kernel_size_list[0] // 2, kernel_size_list[1] // 2
+
+        center_x_pad = [p_vert, -p_vert, 1, p_hor, -p_hor, 1, "center"]
+        upper_x_pad = [0, p_vert, 1, p_hor, -p_hor, 1, "upper"]
+        lower_x_pad = [-p_vert, x_pad.shape[1], 1, p_hor, -p_hor, 1, "lower"]
+        left_x_pad = [p_vert, -p_vert, 1, 0, p_hor, 1, "left"]
+        right_x_pad = [p_vert, -p_vert, 1, -p_hor, x_pad.shape[2], 1, "right"]
+        all_x = [0, x.shape[1], 1, 0, x.shape[2], 1]
+        upper_x = [0, p_vert, 1, 0, x.shape[2], 1]
+        upper_x_rev = [0, p_vert, -1, 0, x.shape[2], 1]
+        upper_x_refl = [1, p_vert + 1, -1, 0, x.shape[2], 1]
+        lower_x = [-p_vert, x.shape[1], 1, 0, x.shape[2], 1]
+        lower_x_rev = [-p_vert, x.shape[1], -1, 0, x.shape[2], 1]
+        lower_x_refl = [-p_vert - 1, x.shape[1] - 1, -1, 0, x.shape[2], 1]
+        left_x = [0, x.shape[1], 1, 0, p_hor, 1]
+        left_x_rev = [0, x.shape[1], 1, 0, p_hor, -1]
+        left_x_refl = [0, x.shape[1], 1, 1, p_hor + 1, -1]
+        right_x = [0, x.shape[1], 1, -p_hor, x.shape[2], 1]
+        right_x_rev = [0, x.shape[1], 1, -p_hor, x.shape[2], -1]
+        right_x_refl = [0, x.shape[1], 1, -p_hor - 1, x.shape[2] - 1, -1]
+        zero_pad = [None, None, None, None]
+        pad_tests = [
+            {
+                "circular": [center_x_pad, all_x],
+                "constant": [center_x_pad, all_x],
+                "symmetric": [center_x_pad, all_x],
+                "reflect": [center_x_pad, all_x],
+            },
+            {
+                "circular": [upper_x_pad, lower_x],
+                "constant": [upper_x_pad, zero_pad],
+                "symmetric": [upper_x_pad, upper_x_rev],
+                "reflect": [upper_x_pad, upper_x_refl],
+            },
+            {
+                "circular": [lower_x_pad, upper_x],
+                "constant": [lower_x_pad, zero_pad],
+                "symmetric": [lower_x_pad, lower_x_rev],
+                "reflect": [lower_x_pad, lower_x_refl],
+            },
+            {
+                "circular": [left_x_pad, right_x],
+                "constant": [left_x_pad, zero_pad],
+                "symmetric": [left_x_pad, left_x_rev],
+                "reflect": [left_x_pad, left_x_refl],
+            },
+            {
+                "circular": [right_x_pad, left_x],
+                "constant": [right_x_pad, zero_pad],
+                "symmetric": [right_x_pad, right_x_rev],
+                "reflect": [right_x_pad, right_x_refl],
+            },
+        ]
+
+        for test_pad in pad_tests:
+            self.compare(
+                x_pad,
+                x,
+                index_x=test_pad[padding_tested][0],
+                index_x_ref=test_pad[padding_tested][1],
+            )
+
+    def _test_predict(self, layer_params, padding_tested, input_shape, batch_size):
+        """Compare predictions between pad+Conv2D and PadConv2D layers."""
+        x = np.random.normal(size=(batch_size,) + input_shape).astype("float32")
+        x_pad = self.pad_input(x, padding_tested, layer_params["kernel_size"])
+        layer_params_ref = layer_params.copy()
+        if padding_tested.lower() == "same":
+            layer_params_ref["padding"] = "same"
+
+        model_ref = generate_k_lip_model(
+            layer_type=tf.keras.layers.Conv2D,
+            layer_params=layer_params_ref,
+            input_shape=x_pad.shape[1:],
+            k=1.0,
+        )
+        y_ref = model_ref.predict(x_pad)
+
+        layer_params_pad = layer_params.copy()
+        layer_params_pad["padding"] = padding_tested
+        model = generate_k_lip_model(
+            layer_type=PadConv2D,
+            layer_params=layer_params_pad,
+            input_shape=input_shape,
+            k=1.0,
+        )
+
+        model.layers[1].kernel.assign(model_ref.layers[1].kernel)
+        if model.layers[1].use_bias:
+            model.layers[1].bias.assign(model_ref.layers[1].bias)
+        y = model.predict(x)
+
+        np.testing.assert_allclose(y_ref, y, 1e-2, 0)
+
+    def _test_vanilla(self, layer_params, padding_tested, input_shape, batch_size):
+        """Compare predictions between PadConv2D and its vanilla export."""
+        x = np.random.normal(size=(batch_size,) + input_shape).astype("float32")
+        layer_params_pad = layer_params.copy()
+        layer_params_pad["padding"] = padding_tested
+        model = generate_k_lip_model(
+            layer_type=PadConv2D,
+            layer_params=layer_params_pad,
+            input_shape=input_shape,
+            k=1.0,
+        )
+        y = model.predict(x)
+
+        model_v = vanillaModel(model)
+        y_v = model_v.predict(x)
+        np.testing.assert_allclose(y_v, y, 1e-2, 0)