RRAEsTorch 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

RRAEsTorch/AE_classes/AE_classes.py

@@ -6,7 +6,6 @@ from RRAEsTorch.utilities import (
     MLP_with_CNN3D_trans,
     stable_SVD,
 )   
-from RRAEsTorch.wrappers import vmap_wrap
 import warnings
 from torch.nn import Linear
 from RRAEsTorch.AE_base import get_autoencoder_base
@@ -42,6 +41,8 @@ def latent_func_strong_RRAE(
     y_approx : jnp.array
         The latent space after the truncation.
     """
+    y = y.T # to get the number of samples in the last dimension, as expected by the SVD function
+
     if apply_basis is not None:
         if get_basis_coeffs:
             return apply_basis, apply_basis.T @ y
@@ -51,7 +52,7 @@ def latent_func_strong_RRAE(
             if get_right_sing:
                 raise ValueError("Can not find right singular vector when projecting on basis")
             return apply_basis.T @ y
-        return apply_basis @ apply_basis.T @ y
+        return (apply_basis @ apply_basis.T @ y).T
 
     k_max = -1 if k_max is None else k_max
 
@@ -97,7 +98,7 @@ def latent_func_strong_RRAE(
         sigs = None
     if ret:
         return u_now, coeffs, sigs
-    return y_approx
+    return y_approx.T
 
 def latent_func_var_strong_RRAE(self, y, k_max=None, epsilon=None, return_dist=False, return_lat_dist=False, **kwargs):
     apply_basis = kwargs.get("apply_basis")
@@ -111,6 +112,9 @@ def latent_func_var_strong_RRAE(self, y, k_max=None, epsilon=None, return_dist=F
         return latent_func_strong_RRAE(self, y, k_max, apply_basis=apply_basis, **kwargs)
 
     basis, coeffs = latent_func_strong_RRAE(self, y, k_max=k_max, get_basis_coeffs=True, apply_basis=apply_basis)
+    
+    coeffs = coeffs.T # to get the number of samples as first dim
+
     if self.typ == "eye":
         mean = coeffs
     elif self.typ == "trainable":
@@ -121,19 +125,19 @@ def latent_func_var_strong_RRAE(self, y, k_max=None, epsilon=None, return_dist=F
     logvar = self.lin_logvar(coeffs)
 
     if return_dist:
-        return mean, logvar
+        return mean.T, logvar.T
 
     std = torch.exp(0.5 * logvar)
     if epsilon is not None:
         if len(epsilon.shape) == 4:
             epsilon = epsilon[0, 0] # to allow tpu sharding
-        z = mean + torch.tensor(epsilon, dtype=torch.float32) * std
+        z = mean + epsilon * std
     else:
         z = mean
 
     if return_lat_dist:
-        return basis @ z, mean, logvar
-    return basis @ z
+        return z @ basis.T, mean.T, logvar.T
+    return z @ basis.T
 
 
 class RRAE_MLP(get_autoencoder_base()):
@@ -506,9 +510,8 @@ class VRRAE_CNN(CNN_Autoencoder):
             count=count,
             **kwargs,
         )
-        v_Linear = vmap_wrap(Linear, -1, count=count)
-        self.lin_mean = v_Linear(k_max, k_max)
-        self.lin_logvar = v_Linear(k_max, k_max)
+        self.lin_mean = Linear(k_max, k_max)
+        self.lin_logvar = Linear(k_max, k_max)
         self.typ = typ
 
     def _perform_in_latent(self, y, *args, k_max=None, epsilon=None, return_dist=False, return_lat_dist=False, **kwargs):
@@ -612,10 +615,6 @@ class VRRAE_CNN1D(CNN1D_Autoencoder):
     typ: int
 
     def __init__(self, channels, input_dim, latent_size, k_max, typ="eye", *, count=1, **kwargs):
-        v_Linear = vmap_wrap(Linear, -1, count=count)
-        self.lin_mean = v_Linear(k_max, k_max,)
-        self.lin_logvar = v_Linear(k_max, k_max)
-        self.typ = typ
         super().__init__(
             channels,
             input_dim,
@@ -623,6 +622,11 @@ class VRRAE_CNN1D(CNN1D_Autoencoder):
             count=count,
             **kwargs,
         )
+        
+        self.lin_mean = Linear(k_max, k_max,)
+        self.lin_logvar = Linear(k_max, k_max)
+        self.typ = typ
+        
 
     def _perform_in_latent(self, y, *args, k_max=None, epsilon=None, return_dist=False, return_lat_dist=False, **kwargs):
         return latent_func_var_strong_RRAE(self, y, k_max, epsilon, return_dist, return_lat_dist, **kwargs)

RRAEsTorch/tests/test_AE_classes_CNN.py

@@ -6,18 +6,12 @@ from RRAEsTorch.AE_classes import (
     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,))
@@ -26,51 +20,51 @@ v_LoRAE_CNN = vmap_wrap(LoRAE_CNN, -1, 1, methods)
 @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 = random.normal(size=(num_samples, channels, width, height))
         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
+        model = RRAE_CNN(
+            x.shape[1], x.shape[2], x.shape[3], latent, num_modes, **kwargs
         )
         y = model.encode(x)
-        assert y.shape == (latent, num_samples)
+        assert y.shape == (num_samples, latent)
         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)
+        assert y.shape == (num_samples, latent)
+        assert model.decode(y).shape == (num_samples, channels, width, height)
 
     def test_Vanilla_CNN(self, latent, num_modes, width, height, channels, num_samples):
-        x = random.normal(size=(channels, width, height, num_samples))
+        x = random.normal(size=(num_samples, channels, width, height))
         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
+        model = Vanilla_AE_CNN(
+            x.shape[1], x.shape[2], x.shape[3], latent, **kwargs
         )
         y = model.encode(x)
-        assert y.shape == (latent, num_samples)
+        assert y.shape == (num_samples, latent)
         x = model.decode(y)
-        assert x.shape == (channels, width, height, num_samples)
+        assert x.shape == (num_samples, channels, width, height)
 
 
     def test_IRMAE_CNN(self, latent, num_modes, width, height, channels, num_samples):
-        x = random.normal(size=(channels, width, height, num_samples))
+        x = random.normal(size=(num_samples, channels, width, height))
         x = torch.tensor(x, dtype=torch.float32)
-        model = v_IRMAE_CNN(
-            x.shape[0], x.shape[1], x.shape[2], latent, linear_l=2
+        model = IRMAE_CNN(
+            x.shape[1], x.shape[2], x.shape[3], latent, linear_l=2
         )
         y = model.encode(x)
-        assert y.shape == (latent, num_samples)
+        assert y.shape == (num_samples, latent)
         assert len(model._encode.layers[-1].layers_l) == 2
         x = model.decode(y)
-        assert x.shape == (channels, width, height, num_samples)
+        assert x.shape == (num_samples, channels, width, height)
 
     def test_LoRAE_CNN(self, latent, num_modes, width, height, channels, num_samples):
-        x = random.normal(size=(channels, width, height, num_samples))
+        x = random.normal(size=(num_samples, channels, width, height))
         x = torch.tensor(x, dtype=torch.float32)
-        model = v_LoRAE_CNN(x.shape[0], x.shape[1], x.shape[2], latent)
+        model = LoRAE_CNN(x.shape[1], x.shape[2], x.shape[3], latent)
         y = model.encode(x)
-        assert y.shape == (latent, num_samples)
+        assert y.shape == (num_samples, latent)
         assert len(model._encode.layers[-1].layers_l) == 1
         x = model.decode(y)
-        assert x.shape == (channels, width, height, num_samples)
+        assert x.shape == (num_samples, channels, width, height)

RRAEsTorch/tests/test_AE_classes_MLP.py

@@ -6,67 +6,59 @@ from RRAEsTorch.AE_classes import (
     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 = random.normal(size=(dim_D, 500))
         x = torch.tensor(x, dtype=torch.float32)
-        model = v_RRAE_MLP(x.shape[0], latent, num_modes)
+        model = RRAE_MLP(x.shape[1], latent, num_modes)
         y = model.encode(x)
-        assert y.shape == (latent, dim_D)
+        assert y.shape == (dim_D, latent)
         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)
+        assert sing_vals[num_modes + 1] < 1e-4
+        assert y.shape == (dim_D, latent)
+        assert model.decode(y).shape == (dim_D, 500)
 
     def test_Vanilla_MLP(self, latent, num_modes, dim_D):
-        x = random.normal(size=(500, dim_D))
+        x = random.normal(size=(dim_D, 500))
         x = torch.tensor(x, dtype=torch.float32)
-        model = v_Vanilla_AE_MLP(x.shape[0], latent)
+        model = Vanilla_AE_MLP(x.shape[1], latent)
         y = model.encode(x)
-        assert y.shape == (latent, dim_D)
+        assert y.shape == (dim_D, latent)
         x = model.decode(y)
-        assert x.shape == (500, dim_D)
+        assert x.shape == (dim_D, 500)
 
     def test_IRMAE_MLP(self, latent, num_modes, dim_D):
-        x = random.normal(size=(500, dim_D))
+        x = random.normal(size=(dim_D, 500))
         x = torch.tensor(x, dtype=torch.float32)
-        model = v_IRMAE_MLP(x.shape[0], latent, linear_l=2)
+        model = IRMAE_MLP(x.shape[1], latent, linear_l=2)
         y = model.encode(x)
-        assert y.shape == (latent, dim_D)
+        assert y.shape == (dim_D, latent)
         assert len(model._encode.layers_l) == 2
         x = model.decode(y)
-        assert x.shape == (500, dim_D)
+        assert x.shape == (dim_D, 500)
 
     def test_LoRAE_MLP(self, latent, num_modes, dim_D):
-        x = random.normal(size=(500, dim_D))
+        x = random.normal(size=(dim_D, 500))
         x = torch.tensor(x, dtype=torch.float32)
-        model = v_LoRAE_MLP(x.shape[0], latent)
+        model = LoRAE_MLP(x.shape[1], latent)
         y = model.encode(x)
-        assert y.shape == (latent, dim_D)
+        assert y.shape == (dim_D, latent)
         assert len(model._encode.layers_l) == 1
         x = model.decode(y)
-        assert x.shape == (500, dim_D)
+        assert x.shape == (dim_D, 500)
 
 def test_getting_SVD_coeffs():
-    data = random.uniform(size=(500, 15))
+    data = random.uniform(size=(15, 500))
     data = torch.tensor(data, dtype=torch.float32)
-    model_s = v_RRAE_MLP(data.shape[0], 200, 3)
+    model_s = RRAE_MLP(data.shape[1], 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

@@ -13,8 +13,8 @@ import torch
 @pytest.mark.parametrize(
     "model_cls, sh, lf",
     [
-        (Vanilla_AE_CNN, (1, 2, 2, 10), "default"),
-        (LoRAE_CNN, (6, 16, 16, 10), "nuc"),
+        (Vanilla_AE_CNN, (10, 1, 2, 2), "default"),
+        (LoRAE_CNN, (10, 6, 16, 16), "nuc"),
     ],
 )
 def test_AE_fitting(model_cls, sh, lf):
@@ -24,10 +24,10 @@ def test_AE_fitting(model_cls, sh, lf):
         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
+        channels=x.shape[1],
+        width=x.shape[2],
+        height=x.shape[3],
+        samples=x.shape[0],  # Only for weak
         k_max=2,
     )
     kwargs = {
@@ -52,16 +52,16 @@ def test_AE_fitting(model_cls, sh, lf):
 def test_IRMAE_fitting():
     model_cls = IRMAE_CNN
     lf = "default"
-    sh = (3, 12, 12, 10)
+    sh = (10, 3, 12, 12)
     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],
+        channels=x.shape[1],
+        width=x.shape[2],
+        height=x.shape[3],
         k_max=2,
         linear_l=4,
     )
@@ -77,7 +77,7 @@ def test_IRMAE_fitting():
         assert False, f"Fitting failed with the following exception {repr(e)}"
 
 def test_RRAE_fitting():
-    sh = (1, 20, 20, 10)
+    sh = (10, 1, 20, 20)
     model_cls = RRAE_CNN
     x = random.normal(size=sh)
     x = torch.tensor(x, dtype=torch.float32)
@@ -85,9 +85,9 @@ def test_RRAE_fitting():
         x,
         model_cls,
         latent_size=100,
-        channels=x.shape[0],
-        width=x.shape[1],
-        height=x.shape[2],
+        channels=x.shape[1],
+        width=x.shape[2],
+        height=x.shape[3],
         k_max=2,
     )
     training_kwargs = {

RRAEsTorch/tests/test_fitting_MLP.py

@@ -14,8 +14,8 @@ import torch
 @pytest.mark.parametrize(
     "model_cls, sh, lf",
     [
-        (Vanilla_AE_MLP, (500, 10), "default"),
-        (LoRAE_MLP, (500, 10), "nuc"),
+        (Vanilla_AE_MLP, (10, 500), "default"),
+        (LoRAE_MLP, (10, 500), "nuc"),
     ],
 )
 def test_fitting(model_cls, sh, lf):
@@ -24,9 +24,8 @@ def test_fitting(model_cls, sh, lf):
     trainor = AE_Trainor_class(
         x,
         model_cls,
-        in_size=x.shape[0],
-        data_size=x.shape[-1],
-        samples=x.shape[-1],  # Only for weak
+        in_size=x.shape[1],
+        samples=x.shape[0],  # Only for weak
         norm_in="meanstd",
         norm_out="minmax",
         out_train=x,
@@ -51,14 +50,14 @@ def test_fitting(model_cls, sh, lf):
 
 
 def test_RRAE_fitting():
-    sh = (500, 10)
+    sh = (10, 500)
     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],
+        in_size=x.shape[1],
         latent_size=2000,
         k_max=2,
     )
@@ -83,14 +82,13 @@ def test_RRAE_fitting():
 def test_IRMAE_fitting():
     model_cls = IRMAE_MLP
     lf = "default"
-    sh = (500, 10)
+    sh = (10, 500)
     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],
+        in_size=x.shape[1],
         latent_size=2000,
         k_max=2,
         linear_l=4,
@@ -107,15 +105,15 @@ def test_IRMAE_fitting():
         assert False, f"Fitting failed with the following exception {repr(e)}"
 
 def test_fitting():
-    sh = (50, 100)
+    sh = (100, 50)
     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]]
+        in_channels=x.shape[1],
+        hidden_channels=[100, x.shape[1]]
     )
     training_kwargs = {
         "step_st": [2],

RRAEsTorch/tests/test_save.py

@@ -5,7 +5,7 @@ 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 = random.normal(size=(1, 1, 28, 28))
     data = torch.tensor(data, dtype=torch.float32)
     model_cls = RRAE_CNN
 
@@ -13,9 +13,9 @@ def test_save():  # Only to test if saving/loading is causing a problem
         data,
         model_cls,
         latent_size=100,
-        channels=data.shape[0],
-        width=data.shape[1],
-        height=data.shape[2],
+        channels=data.shape[1],
+        width=data.shape[2],
+        height=data.shape[3],
         pre_func_inp=lambda x: x * 2 / 17,
         pre_func_out=lambda x: x / 2,
         k_max=2,
@@ -25,17 +25,17 @@ def test_save():  # Only to test if saving/loading is causing a problem
     new_trainor = RRAE_Trainor_class()
     new_trainor.load_model("test_", erase=True)
     try:
-        pr = trainor.model(data[..., 0:1], k_max=2)
+        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)
+        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 = random.normal(size=(1, 1, 28, 28))
     data = torch.tensor(data, dtype=torch.float32)
     
     model_cls = RRAE_CNN
@@ -44,9 +44,9 @@ def test_save_with_final_act():
         data,
         model_cls,
         latent_size=100,
-        channels=data.shape[0],
-        width=data.shape[1],
-        height=data.shape[2],
+        channels=data.shape[1],
+        width=data.shape[2],
+        height=data.shape[3],
         kwargs_dec={"final_activation": torch.sigmoid},
         k_max=2,
     )
@@ -55,7 +55,7 @@ def test_save_with_final_act():
     new_trainor = RRAE_Trainor_class()
     new_trainor.load_model("test_", erase=True)
     try:
-        pr = new_trainor.model(data[..., 0:1], k_max=2)
+        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: