RRAEsTorch 0.1.0__py3-none-any.whl

RRAEsTorch/config.py

@@ -0,0 +1,95 @@
+from RRAEsTorch.utilities import MLP_with_linear   
+import numpy.random as random
+from RRAEsTorch.AE_base import AE_base as rraes
+import torch
+from torch.func import vmap
+
+class Autoencoder(torch.nn.Module):
+    _encode: MLP_with_linear
+    _decode: MLP_with_linear
+    _perform_in_latent: callable
+    _perform_in_latent: callable
+    map_latent: bool
+    norm_funcs: list
+    inv_norm_funcs: list
+    count: int
+
+    def __init__(
+        self,
+        in_size,
+        latent_size,
+        latent_size_after=None,
+        _encode=None,
+        _decode=None,
+        map_latent=True,
+        *,
+        count=1,
+        kwargs_enc={},
+        kwargs_dec={},
+        **kwargs,
+    ):
+        super().__init__()
+        if latent_size_after is None:
+            latent_size_after = latent_size
+
+        if _encode is None:
+            if "width_size" not in kwargs_enc.keys():
+                kwargs_enc["width_size"] = 64
+
+            if "depth" not in kwargs_enc.keys():
+                kwargs_enc["depth"] = 1
+
+            self._encode = MLP_with_linear(
+                in_size=in_size,
+                out_size=latent_size,
+                **kwargs_enc,
+            )
+
+        else:
+            self._encode = _encode
+
+        if not hasattr(self, "_perform_in_latent"):
+            self._perform_in_latent = lambda x, *args, **kwargs: x 
+
+        if _decode is None:
+            if "width_size" not in kwargs_dec.keys():
+                kwargs_dec["width_size"] = 64
+            if "depth" not in kwargs_dec.keys():
+                kwargs_dec["depth"] = 6
+
+            self._decode = MLP_with_linear(
+                in_size=latent_size_after,
+                out_size=in_size,
+                **kwargs_dec,
+            )
+        else:
+            self._decode = _decode
+
+        self.count = count
+        self.map_latent = map_latent
+        self.inv_norm_funcs = ["decode"]
+        self.norm_funcs = ["encode", "latent"]
+
+    def encode(self, x, *args, **kwargs):
+        return self._encode(x, *args, **kwargs)
+    
+    def decode(self, x, *args, **kwargs):
+        return self._decode(x, *args, **kwargs)
+
+    def perform_in_latent(self, y, *args, **kwargs):
+        if self.map_latent:
+            new_perform_in_latent = lambda x: self._perform_in_latent(
+                x, *args, **kwargs
+            )
+            for _ in range(self.count):
+                new_perform_in_latent = vmap(new_perform_in_latent, in_dims=-1, out_dims=-1) 
+            return new_perform_in_latent(y)
+        return self._perform_in_latent(y, *args, **kwargs)
+
+    def forward(self, x, *args, **kwargs):
+        return self.decode(self.perform_in_latent(self.encode(x), *args, **kwargs))
+
+    def latent(self, x, *args, **kwargs):
+        return self.perform_in_latent(self.encode(x), *args, **kwargs)
+
+rraes.set_autoencoder_base(Autoencoder)

RRAEsTorch/tests/test_AE_classes_CNN.py

@@ -0,0 +1,76 @@
+import RRAEsTorch.config
+import pytest
+from RRAEsTorch.AE_classes import (
+    RRAE_CNN,
+    Vanilla_AE_CNN,
+    IRMAE_CNN,
+    LoRAE_CNN,
+)
+from RRAEsTorch.wrappers import vmap_wrap
+import numpy.random as random
+import numpy as np
+import torch
+
+methods = ["encode", "decode"]
+
+v_RRAE_CNN = vmap_wrap(RRAE_CNN, -1, 1, methods)
+v_Vanilla_AE_CNN = vmap_wrap(Vanilla_AE_CNN, -1, 1, methods)
+v_IRMAE_CNN = vmap_wrap(IRMAE_CNN, -1, 1, methods)
+v_LoRAE_CNN = vmap_wrap(LoRAE_CNN, -1, 1, methods)
+
+@pytest.mark.parametrize("width", (10, 17, 149))
+@pytest.mark.parametrize("height", (20,))
+@pytest.mark.parametrize("latent", (200,))
+@pytest.mark.parametrize("num_modes", (1,))
+@pytest.mark.parametrize("channels", (1, 3, 5))
+@pytest.mark.parametrize("num_samples", (10, 100))
+class Test_AEs_shapes:
+    def test_RRAE_CNN(self, latent, num_modes, width, height, channels, num_samples):
+        x = random.normal(size=(channels, width, height, num_samples))
+        x = torch.tensor(x, dtype=torch.float32)
+        kwargs = {"kwargs_dec": {"stride": 2}}
+        model = v_RRAE_CNN(
+            x.shape[0], x.shape[1], x.shape[2], latent, num_modes, **kwargs
+        )
+        y = model.encode(x)
+        assert y.shape == (latent, num_samples)
+        y = model.latent(x, k_max=num_modes)
+        _, sing_vals, _ = torch.linalg.svd(y, full_matrices=False)
+        assert sing_vals[num_modes + 1] < 1e-5
+        assert y.shape == (latent, num_samples)
+        assert model.decode(y).shape == (channels, width, height, num_samples)
+
+    def test_Vanilla_CNN(self, latent, num_modes, width, height, channels, num_samples):
+        x = random.normal(size=(channels, width, height, num_samples))
+        x = torch.tensor(x, dtype=torch.float32)
+        kwargs = {"kwargs_dec": {"stride": 2}}
+        model = v_Vanilla_AE_CNN(
+            x.shape[0], x.shape[1], x.shape[2], latent, **kwargs
+        )
+        y = model.encode(x)
+        assert y.shape == (latent, num_samples)
+        x = model.decode(y)
+        assert x.shape == (channels, width, height, num_samples)
+
+
+    def test_IRMAE_CNN(self, latent, num_modes, width, height, channels, num_samples):
+        x = random.normal(size=(channels, width, height, num_samples))
+        x = torch.tensor(x, dtype=torch.float32)
+        model = v_IRMAE_CNN(
+            x.shape[0], x.shape[1], x.shape[2], latent, linear_l=2
+        )
+        y = model.encode(x)
+        assert y.shape == (latent, num_samples)
+        assert len(model._encode.layers[-1].layers_l) == 2
+        x = model.decode(y)
+        assert x.shape == (channels, width, height, num_samples)
+
+    def test_LoRAE_CNN(self, latent, num_modes, width, height, channels, num_samples):
+        x = random.normal(size=(channels, width, height, num_samples))
+        x = torch.tensor(x, dtype=torch.float32)
+        model = v_LoRAE_CNN(x.shape[0], x.shape[1], x.shape[2], latent)
+        y = model.encode(x)
+        assert y.shape == (latent, num_samples)
+        assert len(model._encode.layers[-1].layers_l) == 1
+        x = model.decode(y)
+        assert x.shape == (channels, width, height, num_samples)

RRAEsTorch/tests/test_AE_classes_MLP.py

@@ -0,0 +1,73 @@
+import RRAEsTorch.config
+import pytest
+from RRAEsTorch.AE_classes import (
+    RRAE_MLP,
+    Vanilla_AE_MLP,
+    IRMAE_MLP,
+    LoRAE_MLP,
+)
+from RRAEsTorch.wrappers import vmap_wrap
+import numpy.random as random
+import numpy as np
+import torch
+
+methods = ["encode", "decode"]
+
+v_RRAE_MLP = vmap_wrap(RRAE_MLP, -1, 1, methods)
+v_Vanilla_AE_MLP = vmap_wrap(Vanilla_AE_MLP, -1, 1, methods)
+v_IRMAE_MLP = vmap_wrap(IRMAE_MLP, -1, 1, methods)
+v_LoRAE_MLP = vmap_wrap(LoRAE_MLP, -1, 1, methods)
+
+@pytest.mark.parametrize("dim_D", (10, 15, 50))
+@pytest.mark.parametrize("latent", (200, 400, 800))
+@pytest.mark.parametrize("num_modes", (1, 2, 6))
+class Test_AEs_shapes:
+    def test_RRAE_MLP(self, latent, num_modes, dim_D):
+        x = random.normal(size=(500, dim_D))
+        x = torch.tensor(x, dtype=torch.float32)
+        model = v_RRAE_MLP(x.shape[0], latent, num_modes)
+        y = model.encode(x)
+        assert y.shape == (latent, dim_D)
+        y = model.perform_in_latent(y, k_max=num_modes)
+        _, sing_vals, _ = torch.linalg.svd(y, full_matrices=False)
+        assert sing_vals[num_modes + 1] < 1e-5
+        assert y.shape == (latent, dim_D)
+        assert model.decode(y).shape == (500, dim_D)
+
+    def test_Vanilla_MLP(self, latent, num_modes, dim_D):
+        x = random.normal(size=(500, dim_D))
+        x = torch.tensor(x, dtype=torch.float32)
+        model = v_Vanilla_AE_MLP(x.shape[0], latent)
+        y = model.encode(x)
+        assert y.shape == (latent, dim_D)
+        x = model.decode(y)
+        assert x.shape == (500, dim_D)
+
+    def test_IRMAE_MLP(self, latent, num_modes, dim_D):
+        x = random.normal(size=(500, dim_D))
+        x = torch.tensor(x, dtype=torch.float32)
+        model = v_IRMAE_MLP(x.shape[0], latent, linear_l=2)
+        y = model.encode(x)
+        assert y.shape == (latent, dim_D)
+        assert len(model._encode.layers_l) == 2
+        x = model.decode(y)
+        assert x.shape == (500, dim_D)
+
+    def test_LoRAE_MLP(self, latent, num_modes, dim_D):
+        x = random.normal(size=(500, dim_D))
+        x = torch.tensor(x, dtype=torch.float32)
+        model = v_LoRAE_MLP(x.shape[0], latent)
+        y = model.encode(x)
+        assert y.shape == (latent, dim_D)
+        assert len(model._encode.layers_l) == 1
+        x = model.decode(y)
+        assert x.shape == (500, dim_D)
+
+def test_getting_SVD_coeffs():
+    data = random.uniform(size=(500, 15))
+    data = torch.tensor(data, dtype=torch.float32)
+    model_s = v_RRAE_MLP(data.shape[0], 200, 3)
+    basis, coeffs = model_s.latent(data, k_max=3, get_basis_coeffs=True)
+    assert basis.shape == (200, 3)
+    assert coeffs.shape == (3, 15)
+

RRAEsTorch/tests/test_fitting_CNN.py

@@ -0,0 +1,109 @@
+import RRAEsTorch.config
+import pytest
+from RRAEsTorch.AE_classes import (
+    RRAE_CNN,
+    Vanilla_AE_CNN,
+    IRMAE_CNN,
+    LoRAE_CNN,
+)
+from RRAEsTorch.training_classes import RRAE_Trainor_class, Trainor_class, AE_Trainor_class
+import numpy.random as random
+import torch
+
+@pytest.mark.parametrize(
+    "model_cls, sh, lf",
+    [
+        (Vanilla_AE_CNN, (1, 2, 2, 10), "default"),
+        (LoRAE_CNN, (6, 16, 16, 10), "nuc"),
+    ],
+)
+def test_AE_fitting(model_cls, sh, lf):
+    x = random.normal(size=sh)
+    x = torch.tensor(x, dtype=torch.float32)
+    trainor = AE_Trainor_class(
+        x,
+        model_cls,
+        latent_size=100,
+        channels=x.shape[0],
+        width=x.shape[1],
+        height=x.shape[2],
+        samples=x.shape[-1],  # Only for weak
+        k_max=2,
+    )
+    kwargs = {
+        "step_st": [2],
+        "loss_kwargs": {
+            "lambda_nuc": 0.001,
+            "find_layer": lambda model: model.encode.layers[-2].layers[-1].weight,
+            "loss_type": lf
+        },
+    }
+    try:
+        trainor.fit(
+            x,
+            x,
+            verbose=False,
+            training_kwargs=kwargs,
+        )
+    except Exception as e:
+        assert False, f"Fitting failed with the following exception {repr(e)}"
+
+
+def test_IRMAE_fitting():
+    model_cls = IRMAE_CNN
+    lf = "default"
+    sh = (3, 12, 12, 10)
+    x = random.normal(size=sh)
+    x = torch.tensor(x, dtype=torch.float32)
+    trainor = AE_Trainor_class(
+        x,
+        model_cls,
+        latent_size=100,
+        channels=x.shape[0],
+        width=x.shape[1],
+        height=x.shape[2],
+        k_max=2,
+        linear_l=4,
+    )
+    kwargs = {"step_st": [2], "loss_type": lf}
+    try:
+        trainor.fit(
+            x,
+            x,
+            verbose=False,
+            training_kwargs=kwargs,
+        )
+    except Exception as e:
+        assert False, f"Fitting failed with the following exception {repr(e)}"
+
+def test_RRAE_fitting():
+    sh = (1, 20, 20, 10)
+    model_cls = RRAE_CNN
+    x = random.normal(size=sh)
+    x = torch.tensor(x, dtype=torch.float32)
+    trainor = RRAE_Trainor_class(
+        x,
+        model_cls,
+        latent_size=100,
+        channels=x.shape[0],
+        width=x.shape[1],
+        height=x.shape[2],
+        k_max=2,
+    )
+    training_kwargs = {
+        "step_st": [2],
+        "loss_type": "default"
+    }
+    ft_kwargs = {
+        "step_st": [2],
+    }
+    try:
+        trainor.fit(
+            x,
+            x,
+            verbose=False,
+            training_kwargs=training_kwargs,
+            ft_kwargs=ft_kwargs,
+        )
+    except Exception as e:
+        assert False, f"Fitting failed with the following exception {repr(e)}"

RRAEsTorch/tests/test_fitting_MLP.py

@@ -0,0 +1,133 @@
+import RRAEsTorch.config
+import numpy.random as random
+import pytest
+from RRAEsTorch.AE_classes import (
+    RRAE_MLP,
+    Vanilla_AE_MLP,
+    IRMAE_MLP,
+    LoRAE_MLP,
+)
+from torchvision.ops import MLP
+from RRAEsTorch.training_classes import RRAE_Trainor_class, Trainor_class, AE_Trainor_class
+import torch
+
+@pytest.mark.parametrize(
+    "model_cls, sh, lf",
+    [
+        (Vanilla_AE_MLP, (500, 10), "default"),
+        (LoRAE_MLP, (500, 10), "nuc"),
+    ],
+)
+def test_fitting(model_cls, sh, lf):
+    x = random.normal(size=sh)
+    x = torch.tensor(x, dtype=torch.float32)
+    trainor = AE_Trainor_class(
+        x,
+        model_cls,
+        in_size=x.shape[0],
+        data_size=x.shape[-1],
+        samples=x.shape[-1],  # Only for weak
+        norm_in="meanstd",
+        norm_out="minmax",
+        out_train=x,
+        latent_size=2000,
+        k_max=2,
+    )
+    kwargs = {
+        "step_st": [2],
+        "loss_kwargs": {"lambda_nuc": 0.001, "find_layer": lambda model: model._encode.layers_l[0].weight},
+        "loss_type": lf
+    }
+    try:
+        trainor.fit(
+            x,
+            x,
+            verbose=False,
+            training_kwargs=kwargs,
+        )
+    except Exception as e:
+        assert False, f"Fitting failed with the following exception {repr(e)}"
+
+
+
+def test_RRAE_fitting():
+    sh = (500, 10)
+    model_cls = RRAE_MLP
+    x = random.normal(size=sh)
+    x = torch.tensor(x, dtype=torch.float32)
+    trainor = RRAE_Trainor_class(
+        x,
+        model_cls,
+        in_size=x.shape[0],
+        latent_size=2000,
+        k_max=2,
+    )
+    training_kwargs = {
+        "step_st": [2],
+        "loss_type":"RRAE"
+    }
+    ft_kwargs = {
+        "step_st": [2],
+    }
+    try:
+        trainor.fit(
+            x,
+            x,
+            verbose=False,
+            training_kwargs=training_kwargs,
+            ft_kwargs=ft_kwargs,    
+        )
+    except Exception as e:
+        assert False, f"Fitting failed with the following exception {repr(e)}"
+
+def test_IRMAE_fitting():
+    model_cls = IRMAE_MLP
+    lf = "default"
+    sh = (500, 10)
+    x = random.normal(size=sh)
+    x = torch.tensor(x, dtype=torch.float32)
+    trainor = AE_Trainor_class(
+        x,
+        model_cls,
+        in_size=x.shape[0],
+        data_size=x.shape[-1],
+        latent_size=2000,
+        k_max=2,
+        linear_l=4,
+    )
+    kwargs = {"step_st": [2], "loss_type":lf}
+    try:
+        trainor.fit(
+            x,
+            x,
+            verbose=False,
+            training_kwargs=kwargs,
+        )
+    except Exception as e:
+        assert False, f"Fitting failed with the following exception {repr(e)}"
+
+def test_fitting():
+    sh = (50, 100)
+    model_cls = MLP
+    x = random.normal(size=sh)
+    x = torch.tensor(x, dtype=torch.float32)
+    trainor = Trainor_class(
+        x,
+        model_cls,
+        in_channels=x.shape[0],
+        hidden_channels=[100, x.shape[0]]
+    )
+    training_kwargs = {
+        "step_st": [2],
+        "loss_type": "default"
+    }
+
+    try:
+        trainor.fit(
+            x,
+            x,
+            verbose=False,
+            **training_kwargs,
+        )
+    except Exception as e:
+        assert False, f"Fitting failed with the following exception {repr(e)}"

RRAEsTorch/tests/test_mains.py

@@ -0,0 +1,34 @@
+import subprocess
+import os
+import pytest
+import shutil
+
+def try_remove(name):
+    try:
+        shutil.rmtree(name)
+    except FileNotFoundError:
+        pass
+
+def run_script(script_name):
+    try:
+        result = subprocess.run(
+            ["python", script_name], check=True, capture_output=True, text=True
+        )
+        try_remove("shift")
+        try_remove("folder_name")
+        try_remove("2d_gaussian_shift_scale")
+        try_remove("gaussian_shift")
+        return result.stdout
+    except subprocess.CalledProcessError as e:
+        pytest.fail(f"Error running {script_name}:\n{e.stderr}")
+
+
+@pytest.mark.parametrize(
+    "script_name", ["main-MLP.py", "main-CNN.py",  "general-MLP.py", "main-adap-CNN.py", "main-adap-MLP.py", "main-var-CNN.py", "main-CNN1D.py"]
+)
+def test_scripts(script_name):
+    if os.path.exists(script_name):
+        output = run_script(script_name)
+        assert output is not None
+    else:
+        pytest.fail(f"Script {script_name} not found")

RRAEsTorch/tests/test_save.py

@@ -0,0 +1,62 @@
+import RRAEsTorch.config
+from RRAEsTorch.AE_classes import RRAE_CNN
+from RRAEsTorch.training_classes import RRAE_Trainor_class
+import numpy.random as random
+import torch
+
+def test_save():  # Only to test if saving/loading is causing a problem
+    data = random.normal(size=(1, 28, 28, 1))
+    data = torch.tensor(data, dtype=torch.float32)
+    model_cls = RRAE_CNN
+
+    trainor = RRAE_Trainor_class(
+        data,
+        model_cls,
+        latent_size=100,
+        channels=data.shape[0],
+        width=data.shape[1],
+        height=data.shape[2],
+        pre_func_inp=lambda x: x * 2 / 17,
+        pre_func_out=lambda x: x / 2,
+        k_max=2,
+    )
+
+    trainor.save_model("test_")
+    new_trainor = RRAE_Trainor_class()
+    new_trainor.load_model("test_", erase=True)
+    try:
+        pr = trainor.model(data[..., 0:1], k_max=2)
+    except Exception as e:
+        raise ValueError(f"Original trainor failed with following exception {e}")
+    try:
+        pr = new_trainor.model(data[..., 0:1], k_max=2)
+    except Exception as e:
+        raise ValueError(f"Failed with following exception {e}")
+
+
+def test_save_with_final_act():
+    data = random.normal(size=(1, 28, 28, 1))
+    data = torch.tensor(data, dtype=torch.float32)
+    
+    model_cls = RRAE_CNN
+
+    trainor = RRAE_Trainor_class(
+        data,
+        model_cls,
+        latent_size=100,
+        channels=data.shape[0],
+        width=data.shape[1],
+        height=data.shape[2],
+        kwargs_dec={"final_activation": torch.sigmoid},
+        k_max=2,
+    )
+
+    trainor.save_model("test_")
+    new_trainor = RRAE_Trainor_class()
+    new_trainor.load_model("test_", erase=True)
+    try:
+        pr = new_trainor.model(data[..., 0:1], k_max=2)
+        assert torch.max(pr) <= 1.0, "Final activation not working"
+        assert torch.min(pr) >= 0.0, "Final activation not working"
+    except Exception as e:
+        raise ValueError(f"Failed with following exception {e}")

RRAEsTorch/tests/test_stable_SVD.py

@@ -0,0 +1,37 @@
+from RRAEsTorch.utilities import stable_SVD
+from torch.linalg import svd as normal_svd
+import numpy.random as random
+import pytest
+import numpy as np
+import torch
+
+def stable_SVD_to_scalar(A):
+    U, s, Vt = stable_SVD(A)
+    return torch.linalg.norm((U * s) @ Vt)  # Any scalar depending on U, s, and Vt.
+
+def normal_svd_to_scalar(A):
+    U, s, Vt = normal_svd(A, full_matrices=False)
+    return torch.linalg.norm((U * s) @ Vt)  # Any scalar depending on U, s, and Vt.
+
+@pytest.mark.parametrize(
+    "length, width",
+    [(10, 10), (100, 10), (10, 100), (50000, 100), (1000, 1000), (100, 50000)],
+)
+def test_random_normal(length, width):
+    A = random.uniform(low=0.0, high=1.0, size=(length, width))
+    A = torch.tensor(A, dtype=torch.float32)
+
+    A = A.clone().detach().requires_grad_(True)
+
+    stable_value = stable_SVD_to_scalar(A)
+    stable_value.backward()
+    stable_grad = A.grad
+
+    A = A.clone().detach().requires_grad_(True)
+
+    normal_value = normal_svd_to_scalar(A)
+    normal_value.backward()
+    normal_grad = A.grad
+    
+    assert torch.allclose(stable_value, normal_value, atol=1e-5, rtol=1e-5)
+    assert torch.allclose(stable_grad, normal_grad, atol=1e-5, rtol=1e-5)

RRAEsTorch/tests/test_wrappers.py

@@ -0,0 +1,56 @@
+from RRAEsTorch.wrappers import vmap_wrap, norm_wrap
+from torchvision.ops import MLP
+import pytest
+import numpy as np
+import math
+import numpy.random as random
+import torch
+
+def test_vmap_wrapper():
+    # Usually MLP only accepts a vector, here we give
+    # a tensor and vectorize over the last axis twice
+    data = random.normal(size=(50, 60, 600))
+    data = torch.tensor(data, dtype=torch.float32)
+    
+    model_cls = vmap_wrap(MLP, -1, 2)
+    model = model_cls(50, [64, 100])
+    try:
+        model(data)
+    except ValueError:
+        pytest.fail("Vmap wrapper is not working properly.")
+
+def test_norm_wrapper():
+    # Testing the keep_normalized kwarg
+    data = random.normal(size=(50,))
+    data = torch.tensor(data, dtype=torch.float32)
+    model_cls = norm_wrap(MLP, data, "minmax", None, data, "minmax", None)
+    model = model_cls(50, [64, 100])
+    try:
+        assert not torch.allclose(model(data), model(data, keep_normalized=True))
+    except AssertionError:
+        pytest.fail("The keep_normalized kwarg for norm wrapper is not behaving as expected.")
+
+    # Testing minmax with knwon mins and maxs
+    data = np.linspace(-1, 1, 100)
+    data = torch.tensor(data, dtype=torch.float32)
+    model_cls = norm_wrap(MLP, data, "minmax", None, data, "minmax", None)
+    model = model_cls(50, [64, 100])
+    try:
+        assert 0.55 == model.norm_in.default(None, 0.1)
+        assert -0.8 == model.inv_norm_out.default(None, 0.1)
+    except AssertionError:
+        pytest.fail("Something wrong with minmax wrapper.")
+
+    # Testing meanstd with knwon mean and std
+    data = random.normal(size=(50,))
+    data = (data-np.mean(data))/np.std(data)
+    data = data*2.0 + 1.0  # mean of 1 and std of 2
+    data = torch.tensor(data, dtype=torch.float32)
+
+    model_cls = norm_wrap(MLP, data, "meanstd", None, data, "meanstd", None)
+    model = model_cls(50, [64, 100])
+    try:
+        assert math.isclose(2, model.norm_in.default(None, 5), rel_tol=1e-1, abs_tol=1e-1)
+        assert math.isclose(7, model.inv_norm_out.default(None, 3), rel_tol=1e-1, abs_tol=1e-1)
+    except AssertionError:
+        pytest.fail("Something wrong with norm wrapper.")

RRAEsTorch/trackers/__init__.py

@@ -0,0 +1 @@
+from .trackers import *