diff --git a/.gitignore b/.gitignore
index b44ea43d7..dd5860a94 100644
--- a/.gitignore
+++ b/.gitignore
@@ -97,6 +97,7 @@ celerybeat-schedule
 # virtualenv
 venv/
 ENV/
+.venv/
 
 # Spyder project settings
 .spyderproject
@@ -120,4 +121,4 @@ debug
 .vscode
 
 # pytest cahche
-.pytest_cache
\ No newline at end of file
+.pytest_cache
diff --git a/RELEASES.md b/RELEASES.md
index 65ad8b35f..df5ade483 100644
--- a/RELEASES.md
+++ b/RELEASES.md
@@ -6,6 +6,8 @@
 + Added support for [Nearest Brenier Potentials (SSNB)](http://proceedings.mlr.press/v108/paty20a/paty20a.pdf) (PR #526)
 + Tweaked `get_backend` to ignore `None` inputs (PR #525)
 + Callbacks for generalized conditional gradient in `ot.da.sinkhorn_l1l2_gl` are now vectorized to improve performance (PR #507)
++ The `linspace` method of the backends now has the `type_as` argument to convert to the same dtype and device. (PR #533)
++ The `convolutional_barycenter2d` and `convolutional_barycenter2d_debiased` functions now work with different devices.. (PR #533)
 
 #### Closed issues
 - Fix line search evaluating cost outside of the interpolation range (Issue #502, PR #504)
diff --git a/ot/backend.py b/ot/backend.py
index e9750ee0a..80eb7604e 100644
--- a/ot/backend.py
+++ b/ot/backend.py
@@ -86,15 +86,15 @@
 #
 # License: MIT License
 
-import numpy as np
 import os
-import scipy
-import scipy.linalg
-from scipy.sparse import issparse, coo_matrix, csr_matrix
-import scipy.special as special
 import time
 import warnings
 
+import numpy as np
+import scipy
+import scipy.linalg
+import scipy.special as special
+from scipy.sparse import coo_matrix, csr_matrix, issparse
 
 DISABLE_TORCH_KEY = 'POT_BACKEND_DISABLE_PYTORCH'
 DISABLE_JAX_KEY = 'POT_BACKEND_DISABLE_JAX'
@@ -650,7 +650,7 @@ def std(self, a, axis=None):
         """
         raise NotImplementedError()
 
-    def linspace(self, start, stop, num):
+    def linspace(self, start, stop, num, type_as=None):
         r"""
         Returns a specified number of evenly spaced values over a given interval.
 
@@ -1208,8 +1208,11 @@ def median(self, a, axis=None):
     def std(self, a, axis=None):
         return np.std(a, axis=axis)
 
-    def linspace(self, start, stop, num):
-        return np.linspace(start, stop, num)
+    def linspace(self, start, stop, num, type_as=None):
+        if type_as is None:
+            return np.linspace(start, stop, num)
+        else:
+            return np.linspace(start, stop, num, dtype=type_as.dtype)
 
     def meshgrid(self, a, b):
         return np.meshgrid(a, b)
@@ -1579,8 +1582,11 @@ def median(self, a, axis=None):
     def std(self, a, axis=None):
         return jnp.std(a, axis=axis)
 
-    def linspace(self, start, stop, num):
-        return jnp.linspace(start, stop, num)
+    def linspace(self, start, stop, num, type_as=None):
+        if type_as is None:
+            return jnp.linspace(start, stop, num)
+        else:
+            return self._change_device(jnp.linspace(start, stop, num, dtype=type_as.dtype), type_as)
 
     def meshgrid(self, a, b):
         return jnp.meshgrid(a, b)
@@ -1986,6 +1992,7 @@ def concatenate(self, arrays, axis=0):
 
     def zero_pad(self, a, pad_width, value=0):
         from torch.nn.functional import pad
+
         # pad_width is an array of ndim tuples indicating how many 0 before and after
         # we need to add. We first need to make it compliant with torch syntax, that
         # starts with the last dim, then second last, etc.
@@ -2006,6 +2013,7 @@ def mean(self, a, axis=None):
 
     def median(self, a, axis=None):
         from packaging import version
+
         # Since version 1.11.0, interpolation is available
         if version.parse(torch.__version__) >= version.parse("1.11.0"):
             if axis is not None:
@@ -2026,8 +2034,11 @@ def std(self, a, axis=None):
         else:
             return torch.std(a, unbiased=False)
 
-    def linspace(self, start, stop, num):
-        return torch.linspace(start, stop, num, dtype=torch.float64)
+    def linspace(self, start, stop, num, type_as=None):
+        if type_as is None:
+            return torch.linspace(start, stop, num)
+        else:
+            return torch.linspace(start, stop, num, dtype=type_as.dtype, device=type_as.device)
 
     def meshgrid(self, a, b):
         try:
@@ -2427,8 +2438,12 @@ def median(self, a, axis=None):
     def std(self, a, axis=None):
         return cp.std(a, axis=axis)
 
-    def linspace(self, start, stop, num):
-        return cp.linspace(start, stop, num)
+    def linspace(self, start, stop, num, type_as=None):
+        if type_as is None:
+            return cp.linspace(start, stop, num)
+        else:
+            with cp.cuda.Device(type_as.device):
+                return cp.linspace(start, stop, num, dtype=type_as.dtype)
 
     def meshgrid(self, a, b):
         return cp.meshgrid(a, b)
@@ -2834,8 +2849,11 @@ def median(self, a, axis=None):
     def std(self, a, axis=None):
         return tnp.std(a, axis=axis)
 
-    def linspace(self, start, stop, num):
-        return tnp.linspace(start, stop, num)
+    def linspace(self, start, stop, num, type_as=None):
+        if type_as is None:
+            return tnp.linspace(start, stop, num)
+        else:
+            return tnp.linspace(start, stop, num, dtype=type_as.dtype)
 
     def meshgrid(self, a, b):
         return tnp.meshgrid(a, b)
diff --git a/ot/bregman.py b/ot/bregman.py
index 29bcd5867..c90d89986 100644
--- a/ot/bregman.py
+++ b/ot/bregman.py
@@ -20,7 +20,8 @@
 import numpy as np
 from scipy.optimize import fmin_l_bfgs_b
 
-from ot.utils import unif, dist, list_to_array
+from ot.utils import dist, list_to_array, unif
+
 from .backend import get_backend
 
 
@@ -2217,11 +2218,11 @@ def _convolutional_barycenter2d(A, reg, weights=None, numItermax=10000,
 
     # build the convolution operator
     # this is equivalent to blurring on horizontal then vertical directions
-    t = nx.linspace(0, 1, A.shape[1])
+    t = nx.linspace(0, 1, A.shape[1], type_as=A)
     [Y, X] = nx.meshgrid(t, t)
     K1 = nx.exp(-(X - Y) ** 2 / reg)
 
-    t = nx.linspace(0, 1, A.shape[2])
+    t = nx.linspace(0, 1, A.shape[2], type_as=A)
     [Y, X] = nx.meshgrid(t, t)
     K2 = nx.exp(-(X - Y) ** 2 / reg)
 
@@ -2295,11 +2296,11 @@ def _convolutional_barycenter2d_log(A, reg, weights=None, numItermax=10000,
     err = 1
     # build the convolution operator
     # this is equivalent to blurring on horizontal then vertical directions
-    t = nx.linspace(0, 1, width)
+    t = nx.linspace(0, 1, width, type_as=A)
     [Y, X] = nx.meshgrid(t, t)
     M1 = - (X - Y) ** 2 / reg
 
-    t = nx.linspace(0, 1, height)
+    t = nx.linspace(0, 1, height, type_as=A)
     [Y, X] = nx.meshgrid(t, t)
     M2 = - (X - Y) ** 2 / reg
 
@@ -2452,11 +2453,11 @@ def _convolutional_barycenter2d_debiased(A, reg, weights=None, numItermax=10000,
 
     # build the convolution operator
     # this is equivalent to blurring on horizontal then vertical directions
-    t = nx.linspace(0, 1, width)
+    t = nx.linspace(0, 1, width, type_as=A)
     [Y, X] = nx.meshgrid(t, t)
     K1 = nx.exp(-(X - Y) ** 2 / reg)
 
-    t = nx.linspace(0, 1, height)
+    t = nx.linspace(0, 1, height, type_as=A)
     [Y, X] = nx.meshgrid(t, t)
     K2 = nx.exp(-(X - Y) ** 2 / reg)
 
@@ -2532,11 +2533,11 @@ def _convolutional_barycenter2d_debiased_log(A, reg, weights=None, numItermax=10
     err = 1
     # build the convolution operator
     # this is equivalent to blurring on horizontal then vertical directions
-    t = nx.linspace(0, 1, width)
+    t = nx.linspace(0, 1, width, type_as=A)
     [Y, X] = nx.meshgrid(t, t)
     M1 = - (X - Y) ** 2 / reg
 
-    t = nx.linspace(0, 1, height)
+    t = nx.linspace(0, 1, height, type_as=A)
     [Y, X] = nx.meshgrid(t, t)
     M2 = - (X - Y) ** 2 / reg
 
diff --git a/test/test_backend.py b/test/test_backend.py
index 8ab861078..605e30ad8 100644
--- a/test/test_backend.py
+++ b/test/test_backend.py
@@ -6,16 +6,13 @@
 #
 # License: MIT License
 
-import ot
-import ot.backend
-from ot.backend import torch, jax, tf
-
-import pytest
-
 import numpy as np
+import pytest
 from numpy.testing import assert_array_almost_equal_nulp
 
-from ot.backend import get_backend, get_backend_list, to_numpy
+import ot
+import ot.backend
+from ot.backend import get_backend, get_backend_list, jax, tf, to_numpy, torch
 
 
 def test_get_backend_list():
@@ -507,6 +504,7 @@ def test_func_backends(nx):
         lst_name.append('std')
 
         A = nx.linspace(0, 1, 50)
+        A = nx.linspace(0, 1, 50, type_as=Mb)
         lst_b.append(nx.to_numpy(A))
         lst_name.append('linspace')
 
diff --git a/test/test_bregman.py b/test/test_bregman.py
index 8355cda95..8627df3c6 100644
--- a/test/test_bregman.py
+++ b/test/test_bregman.py
@@ -14,7 +14,7 @@
 import pytest
 
 import ot
-from ot.backend import torch, tf
+from ot.backend import tf, torch
 
 
 @pytest.mark.parametrize("verbose, warn", product([True, False], [True, False]))
@@ -352,13 +352,15 @@ def test_sinkhorn2_variants_device_tf(method):
         nx.assert_same_dtype_device(Mb, Gb)
         nx.assert_same_dtype_device(Mb, lossb)
 
+    # Check that everything happens on the GPU
+    ub, Mb = nx.from_numpy(u, M)
+    Gb = ot.sinkhorn(ub, ub, Mb, 1, method=method, stopThr=1e-10)
+    lossb = ot.sinkhorn2(ub, ub, Mb, 1, method=method, stopThr=1e-10)
+    nx.assert_same_dtype_device(Mb, Gb)
+    nx.assert_same_dtype_device(Mb, lossb)
+
+    # Check this only if GPU is available
     if len(tf.config.list_physical_devices('GPU')) > 0:
-        # Check that everything happens on the GPU
-        ub, Mb = nx.from_numpy(u, M)
-        Gb = ot.sinkhorn(ub, ub, Mb, 1, method=method, stopThr=1e-10)
-        lossb = ot.sinkhorn2(ub, ub, Mb, 1, method=method, stopThr=1e-10)
-        nx.assert_same_dtype_device(Mb, Gb)
-        nx.assert_same_dtype_device(Mb, lossb)
         assert nx.dtype_device(Gb)[1].startswith("GPU")
 
 
@@ -691,21 +693,33 @@ def test_barycenter_stabilization(nx):
     np.testing.assert_allclose(bar, bar_np)
 
 
-@pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_log"])
-def test_wasserstein_bary_2d(nx, method):
-    rng = np.random.RandomState(42)
-    size = 20  # size of a square image
+def create_random_images_dist(seed, size=20):
+    """Creates an array of two random images of size (size, size). Returns an array of shape (2, size, size)."""
+    rng = np.random.RandomState(seed)
+
+    # First image
     a1 = rng.rand(size, size)
     a1 += a1.min()
-    a1 = a1 / np.sum(a1)
+    a1 = a1 / np.sum(a1)  # Ensure that it is a probability distribution
+
+    # Second image
     a2 = rng.rand(size, size)
     a2 += a2.min()
     a2 = a2 / np.sum(a2)
-    # creating matrix A containing all distributions
+
+    # Creating matrix A containing all distributions
     A = np.zeros((2, size, size))
     A[0, :, :] = a1
     A[1, :, :] = a2
 
+    return A
+
+
+@pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_log"])
+def test_wasserstein_bary_2d(nx, method):
+    # Create the array of images to test
+    A = create_random_images_dist(42, size=20)
+
     A_nx = nx.from_numpy(A)
 
     # wasserstein
@@ -715,9 +729,108 @@ def test_wasserstein_bary_2d(nx, method):
             ot.bregman.convolutional_barycenter2d(A_nx, reg, method=method)
     else:
         bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d(
-            A, reg, method=method, verbose=True, log=True)
+            A, reg, method=method, verbose=True, log=True
+        )
         bary_wass = nx.to_numpy(
-            ot.bregman.convolutional_barycenter2d(A_nx, reg, method=method))
+            ot.bregman.convolutional_barycenter2d(A_nx, reg, method=method)
+        )
+
+        np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3)
+        np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3)
+
+        # help in checking if log and verbose do not bug the function
+        ot.bregman.convolutional_barycenter2d(A, reg, log=True, verbose=True)
+
+
+@pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_log"])
+def test_wasserstein_bary_2d_dtype_device(nx, method):
+    # Create the array of images to test
+    A = create_random_images_dist(42, size=20)
+
+    for tp in nx.__type_list__:
+        print(nx.dtype_device(tp))
+
+        Ab = nx.from_numpy(A, type_as=tp)
+
+        # wasserstein
+        reg = 1e-2
+        if nx.__name__ in ("jax", "tf") and method == "sinkhorn_log":
+            with pytest.raises(NotImplementedError):
+                ot.bregman.convolutional_barycenter2d(Ab, reg, method=method)
+        else:
+            # Compute the barycenter with numpy
+            bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d(
+                A, reg, method=method, verbose=True, log=True
+            )
+            # Compute the barycenter with the backend
+            bary_wass_b = ot.bregman.convolutional_barycenter2d(Ab, reg, method=method)
+            # Convert the backend result to numpy, to compare with the numpy result
+            bary_wass = nx.to_numpy(bary_wass_b)
+
+            np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3)
+            np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3)
+
+            # help in checking if log and verbose do not bug the function
+            ot.bregman.convolutional_barycenter2d(A, reg, log=True, verbose=True)
+
+            # Test that the dtype and device are the same after the computation
+            nx.assert_same_dtype_device(Ab, bary_wass_b)
+
+
+@pytest.mark.skipif(not tf, reason="tf not installed")
+@pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_log"])
+def test_wasserstein_bary_2d_device_tf(method):
+    # Using the Tensorflow backend
+    nx = ot.backend.TensorflowBackend()
+
+    # Create the array of images to test
+    A = create_random_images_dist(42, size=20)
+
+    # Check that everything stays on the CPU
+    with tf.device("/CPU:0"):
+        Ab = nx.from_numpy(A)
+
+        # wasserstein
+        reg = 1e-2
+        if method == "sinkhorn_log":
+            with pytest.raises(NotImplementedError):
+                ot.bregman.convolutional_barycenter2d(Ab, reg, method=method)
+        else:
+            # Compute the barycenter with numpy
+            bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d(
+                A, reg, method=method, verbose=True, log=True
+            )
+            # Compute the barycenter with the backend
+            bary_wass_b = ot.bregman.convolutional_barycenter2d(Ab, reg, method=method)
+            # Convert the backend result to numpy, to compare with the numpy result
+            bary_wass = nx.to_numpy(bary_wass_b)
+
+            np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3)
+            np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3)
+
+            # help in checking if log and verbose do not bug the function
+            ot.bregman.convolutional_barycenter2d(A, reg, log=True, verbose=True)
+
+            # Test that the dtype and device are the same after the computation
+            nx.assert_same_dtype_device(Ab, bary_wass_b)
+
+    # Check that everything happens on the GPU
+    Ab = nx.from_numpy(A)
+
+    # wasserstein
+    reg = 1e-2
+    if method == "sinkhorn_log":
+        with pytest.raises(NotImplementedError):
+            ot.bregman.convolutional_barycenter2d(Ab, reg, method=method)
+    else:
+        # Compute the barycenter with numpy
+        bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d(
+            A, reg, method=method, verbose=True, log=True
+        )
+        # Compute the barycenter with the backend
+        bary_wass_b = ot.bregman.convolutional_barycenter2d(Ab, reg, method=method)
+        # Convert the backend result to numpy, to compare with the numpy result
+        bary_wass = nx.to_numpy(bary_wass_b)
 
         np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3)
         np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3)
@@ -725,21 +838,18 @@ def test_wasserstein_bary_2d(nx, method):
         # help in checking if log and verbose do not bug the function
         ot.bregman.convolutional_barycenter2d(A, reg, log=True, verbose=True)
 
+        # Test that the dtype and device are the same after the computation
+        nx.assert_same_dtype_device(Ab, bary_wass_b)
+
+        # Check this only if GPU is available
+        if len(tf.config.list_physical_devices("GPU")) > 0:
+            assert nx.dtype_device(bary_wass_b)[1].startswith("GPU")
+
 
 @pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_log"])
 def test_wasserstein_bary_2d_debiased(nx, method):
-    rng = np.random.RandomState(42)
-    size = 20  # size of a square image
-    a1 = rng.rand(size, size)
-    a1 += a1.min()
-    a1 = a1 / np.sum(a1)
-    a2 = rng.rand(size, size)
-    a2 += a2.min()
-    a2 = a2 / np.sum(a2)
-    # creating matrix A containing all distributions
-    A = np.zeros((2, size, size))
-    A[0, :, :] = a1
-    A[1, :, :] = a2
+    # Create the array of images to test
+    A = create_random_images_dist(42, size=20)
 
     A_nx = nx.from_numpy(A)
 
@@ -747,19 +857,123 @@ def test_wasserstein_bary_2d_debiased(nx, method):
     reg = 1e-2
     if nx.__name__ in ("jax", "tf") and method == "sinkhorn_log":
         with pytest.raises(NotImplementedError):
-            ot.bregman.convolutional_barycenter2d_debiased(
-                A_nx, reg, method=method)
+            ot.bregman.convolutional_barycenter2d_debiased(A_nx, reg, method=method)
     else:
         bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d_debiased(
-            A, reg, method=method, verbose=True, log=True)
+            A, reg, method=method, verbose=True, log=True
+        )
         bary_wass = nx.to_numpy(
-            ot.bregman.convolutional_barycenter2d_debiased(A_nx, reg, method=method))
+            ot.bregman.convolutional_barycenter2d_debiased(A_nx, reg, method=method)
+        )
 
         np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3)
         np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3)
 
         # help in checking if log and verbose do not bug the function
-        ot.bregman.convolutional_barycenter2d(A, reg, log=True, verbose=True)
+        ot.bregman.convolutional_barycenter2d_debiased(A, reg, log=True, verbose=True)
+
+
+@pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_log"])
+def test_wasserstein_bary_2d_debiased_dtype_device(nx, method):
+    # Create the array of images to test
+    A = create_random_images_dist(42, size=20)
+
+    for tp in nx.__type_list__:
+        print(nx.dtype_device(tp))
+
+        Ab = nx.from_numpy(A, type_as=tp)
+
+        # wasserstein
+        reg = 1e-2
+        if nx.__name__ in ("jax", "tf") and method == "sinkhorn_log":
+            with pytest.raises(NotImplementedError):
+                ot.bregman.convolutional_barycenter2d_debiased(Ab, reg, method=method)
+        else:
+            # Compute the barycenter with numpy
+            bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d_debiased(
+                A, reg, method=method, verbose=True, log=True
+            )
+            # Compute the barycenter with the backend
+            bary_wass_b = ot.bregman.convolutional_barycenter2d_debiased(
+                Ab, reg, method=method
+            )
+            # Convert the backend result to numpy, to compare with the numpy result
+            bary_wass = nx.to_numpy(bary_wass_b)
+
+            np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3)
+            np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3)
+
+            # help in checking if log and verbose do not bug the function
+            ot.bregman.convolutional_barycenter2d_debiased(
+                A, reg, log=True, verbose=True
+            )
+
+            # Test that the dtype and device are the same after the computation
+            nx.assert_same_dtype_device(Ab, bary_wass_b)
+
+
+@pytest.mark.skipif(not tf, reason="tf not installed")
+@pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_log"])
+def test_wasserstein_bary_2d_debiased_device_tf(method):
+    # Using the Tensorflow backend
+    nx = ot.backend.TensorflowBackend()
+
+    # Create the array of images to test
+    A = create_random_images_dist(42, size=20)
+
+    # Check that everything stays on the CPU
+    with tf.device("/CPU:0"):
+        Ab = nx.from_numpy(A)
+
+        # wasserstein
+        reg = 1e-2
+        if method == "sinkhorn_log":
+            with pytest.raises(NotImplementedError):
+                ot.bregman.convolutional_barycenter2d_debiased(Ab, reg, method=method)
+        else:
+            # Compute the barycenter with numpy
+            bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d_debiased(
+                A, reg, method=method, verbose=True, log=True
+            )
+            # Compute the barycenter with the backend
+            bary_wass_b = ot.bregman.convolutional_barycenter2d_debiased(
+                Ab, reg, method=method
+            )
+            # Convert the backend result to numpy, to compare with the numpy result
+            bary_wass = nx.to_numpy(bary_wass_b)
+
+            np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3)
+            np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3)
+
+            # help in checking if log and verbose do not bug the function
+            ot.bregman.convolutional_barycenter2d_debiased(
+                A, reg, log=True, verbose=True
+            )
+
+            # Test that the dtype and device are the same after the computation
+            nx.assert_same_dtype_device(Ab, bary_wass_b)
+
+    # Check that everything happens on the GPU
+    Ab = nx.from_numpy(A)
+
+    # wasserstein
+    reg = 1e-2
+    if method == "sinkhorn_log":
+        with pytest.raises(NotImplementedError):
+            ot.bregman.convolutional_barycenter2d_debiased(Ab, reg, method=method)
+    else:
+        # Compute the barycenter with numpy
+        bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d_debiased(
+            A, reg, method=method, verbose=True, log=True
+        )
+        # Compute the barycenter with the backend
+        bary_wass_b = ot.bregman.convolutional_barycenter2d_debiased(
+            Ab, reg, method=method
+        )
+        # Convert the backend result to numpy, to compare with the numpy result
+        bary_wass = nx.to_numpy(bary_wass_b)
+
+        np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3)
 
 
 def test_unmix(nx):