RR-layer 0.1.1__py3-none-any.whl

RR_layer/__init__.py

@@ -0,0 +1 @@
+from .rr_layer import RRLayer

RR_layer/rr_layer.py

@@ -0,0 +1,375 @@
+import warnings
+import torch
+import torch.nn as nn
+from collections import deque
+
+def stable_SVD(A):
+    """ Computes a numerically stable SVD.
+    
+    A: input matrix (..., m, n)
+    """
+    return StableSVD.apply(A)
+
+class StableSVD(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, A):
+        # Just compute standard SVD
+        U, S, Vh = torch.linalg.svd(A, full_matrices=False)
+
+        # Save for backward
+        ctx.save_for_backward(U, S, Vh)
+        ctx.original_shape = A.shape
+
+        return U, S, Vh
+
+    @staticmethod
+    def backward(ctx, dU, dS, dVh):
+        """
+        Backward pass for stable SVD.
+        Computes gradient w.r.t. input A given gradients on U, S, Vh.
+        """
+        U, S, Vh = ctx.saved_tensors
+        m, n = ctx.original_shape[-2:]
+
+        dtype = U.dtype
+        device = U.device
+
+        H = lambda x: x.transpose(-2, -1).conj()
+        T = lambda x: x.transpose(-2, -1)
+
+        # Diagonal helpers
+        def diag_embed(x):
+            return torch.diag_embed(x)
+
+        # Singular value vector to diagonal for broadcasting
+        S_mat = S.unsqueeze(-2)
+        S_diff = S_mat - S_mat.transpose(-2, -1)
+
+        # F matrix for repeated singular values
+        eps = 1e-20
+        F = torch.where(S_diff.abs() > eps, 1.0 / S_diff, torch.zeros_like(S_diff))
+
+        # Gradient from singular values
+        dA = U @ diag_embed(dS) @ Vh
+
+        # Contributions from U
+        Ut_dU = H(U) @ dU
+        skew_U = F * (Ut_dU - Ut_dU.transpose(-2, -1))
+        dA += U @ skew_U @ diag_embed(S) @ Vh
+
+        # Contributions from V
+        V = H(Vh)  # n x k
+        Vt_dV = H(V) @ H(dVh)  # k x k
+        skew_V = F * (Vt_dV - Vt_dV.transpose(-2, -1))
+        dA += U @ diag_embed(S) @ skew_V @ Vh
+
+        # Rectangular adjustments (like JAX)
+        # s_inv = 1 / S with stable zero handling
+        s_zeros = (S == 0).to(dtype)
+        s_inv = 1.0 / (S + s_zeros) - s_zeros  # shape: (k,)
+
+        if m > n:
+            dAV = dA @ V
+            dA += (dAV - U @ (H(U) @ dAV)) * s_inv
+        elif n > m:
+            dAHU = H(dA) @ U
+            print("first term shape: ", H(dAHU - V @ (Vh @ dAHU)).shape)
+            print("s_inv shape: ", s_inv.shape)
+
+            dA += H(dAHU - V @ (Vh @ dAHU)) * s_inv.unsqueeze(1)
+
+        return dA
+
+
+import math
+import warnings
+from collections import deque
+
+import torch
+import torch.nn as nn
+
+
+class RRLayer(nn.Module):
+    r"""
+    Rank Reduction (RR) layer.
+
+    During training, the layer computes a truncated SVD across the batch
+    dimension and reconstructs the input using the top ``rank`` singular
+    components.
+
+    During evaluation, the layer projects the input onto a learned inference
+    basis obtained from the recent training bases.
+
+    A custom basis can always be supplied through the ``basis`` argument of
+    :meth:`forward`, in which case the train/eval behavior is bypassed.
+
+    Args:
+        rank (int):
+            Number of singular values to retain.
+
+        basis_history_size (int, optional):
+            Number of recent batch bases to keep when estimating the inference
+            basis. Default: 20.
+
+    Shape:
+        - Input: ``(N, *)``
+        - Output: ``(N, *)``
+
+    Example:
+        >>> rr = RRLayer(rank=8)
+        >>> rr.train()
+        >>> y = rr(torch.randn(32, 768))
+
+        >>> rr.eval()
+        >>> y = rr(torch.randn(32, 768))
+    """
+
+    def __init__(
+        self,
+        rank: int,
+        basis_history_size: int = 20,
+    ):
+        super().__init__()
+
+        if rank <= 0:
+            raise ValueError(
+                f"rank must be positive, got {rank}."
+            )
+
+        if basis_history_size <= 0:
+            raise ValueError(
+                f"basis_history_size must be positive, got "
+                f"{basis_history_size}."
+            )
+
+        self.rank = rank
+        self.basis_history_size = basis_history_size
+
+        self.register_buffer(
+            "inference_basis",
+            None,
+            persistent=True,
+        )
+
+        self._basis_bank = deque(
+            maxlen=basis_history_size,
+        )
+
+    def extra_repr(self) -> str:
+        return (
+            f"rank={self.rank}, "
+            f"basis_history_size={self.basis_history_size}"
+        )
+
+    def train(self, mode: bool = True):
+        """
+        Switch between training and evaluation mode.
+
+        When switching from training to evaluation for the first time,
+        an inference basis is automatically computed from the stored
+        training bases.
+        """
+
+        previous_mode = self.training
+
+        super().train(mode)
+
+        if (
+            previous_mode
+            and not mode
+            and self.inference_basis is None
+            and len(self._basis_bank) > 0
+        ):
+            self.finalize_basis()
+
+        return self
+
+    @torch.no_grad()
+    def finalize_basis(self) -> None:
+        """
+        Build the inference basis from the stored training bases.
+        """
+
+        if len(self._basis_bank) == 0:
+            raise RuntimeError(
+                "Cannot finalize basis: no stored bases available."
+            )
+
+        device = next(self.parameters(), None)
+
+        if device is None:
+            device = self.inference_basis.device \
+                if self.inference_basis is not None \
+                else self._basis_bank[0].device
+        else:
+            device = device.device
+
+        W = torch.cat(
+            [
+                basis.to(device)
+                for basis in self._basis_bank
+            ],
+            dim=1,
+        )
+
+        U, _, _ = stable_SVD(W)
+
+        r = min(self.rank, U.shape[1])
+
+        self.inference_basis = U[:, :r]
+
+    def _validate_inputs(
+        self,
+        x: torch.Tensor,
+        basis: torch.Tensor | None,
+    ) -> None:
+        """
+        Validate inputs for RRLayer.
+        """
+
+        if not isinstance(x, torch.Tensor):
+            raise TypeError(
+                f"x must be a torch.Tensor, got {type(x)}."
+            )
+
+        if x.ndim < 2:
+            raise ValueError(
+                "RRLayer expects input of shape (N, ...), "
+                f"got shape {tuple(x.shape)}."
+            )
+
+        if x.shape[0] == 0:
+            raise ValueError(
+                "Batch size must be greater than zero."
+            )
+
+        if not torch.isfinite(x).all():
+            raise ValueError(
+                "Input tensor contains NaN or Inf values."
+            )
+
+        M = math.prod(x.shape[1:])
+        N = x.shape[0]
+
+        rank_max = min(M, N)
+
+        if self.rank > rank_max:
+            warnings.warn(
+                f"Requested rank={self.rank}, but the maximum "
+                f"achievable rank is {rank_max}. "
+                f"Using rank={rank_max}.",
+                stacklevel=2,
+            )
+
+        if basis is not None:
+
+            if not isinstance(basis, torch.Tensor):
+                raise TypeError(
+                    f"basis must be a torch.Tensor, got {type(basis)}."
+                )
+
+            if basis.ndim != 2:
+                raise ValueError(
+                    "basis must have shape (M, r), "
+                    f"got shape {tuple(basis.shape)}."
+                )
+
+            if basis.shape[0] != M:
+                raise ValueError(
+                    f"basis has {basis.shape[0]} rows but "
+                    f"expected {M}."
+                )
+
+            if basis.shape[1] == 0:
+                raise ValueError(
+                    "basis must contain at least one column."
+                )
+
+            if basis.device != x.device:
+                raise ValueError(
+                    f"basis is on {basis.device} while "
+                    f"x is on {x.device}."
+                )
+
+            if basis.dtype != x.dtype:
+                raise ValueError(
+                    f"basis dtype ({basis.dtype}) does not match "
+                    f"x dtype ({x.dtype})."
+                )
+
+            if not torch.isfinite(basis).all():
+                raise ValueError(
+                    "basis contains NaN or Inf values."
+                )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        basis: torch.Tensor | None = None,
+        return_factors: bool = False,
+    ):
+        self._validate_inputs(x, basis)
+
+        original_shape = x.shape
+        batch_size = original_shape[0]
+
+        X = torch.movedim(x, 0, -1)
+        X = X.reshape(-1, batch_size)
+
+        if basis is not None:
+
+            basis_used = basis
+
+            coeffs = basis_used.T @ X
+
+            X_hat = basis_used @ coeffs
+
+        elif self.training:
+
+            U, S, Vh = stable_SVD(X)
+
+            r = min(self.rank, S.shape[0])
+
+            basis_used = U[:, :r]
+
+            coeffs = (
+                S[:r].unsqueeze(1)
+                * Vh[:r]
+            )
+
+            X_hat = basis_used @ coeffs
+
+            self._basis_bank.append(
+                basis_used.detach().cpu()
+            )
+
+        else:
+
+            if self.inference_basis is None:
+                raise RuntimeError(
+                    "RRLayer has no inference basis. "
+                    "Call finalize_basis() or run training "
+                    "before evaluation."
+                )
+
+            basis_used = self.inference_basis
+
+            coeffs = basis_used.T @ X
+
+            X_hat = basis_used @ coeffs
+
+        output = X_hat.reshape(
+            *original_shape[1:],
+            batch_size,
+        )
+
+        output = torch.movedim(
+            output,
+            -1,
+            0,
+        )
+
+        if return_factors:
+            return output, basis_used, coeffs
+
+        return output

rr_layer-0.1.1.dist-info/METADATA

@@ -0,0 +1,100 @@
+Metadata-Version: 2.4
+Name: RR_layer
+Version: 0.1.1
+Summary:  A repo for RR layer in Pytorch.
+Author-email: Jad Mounayer <jad.mounayer@outlook.com>
+License: MIT
+License-File: LICENSE
+Requires-Python: >=3.10
+Requires-Dist: matplotlib
+Requires-Dist: numpy
+Requires-Dist: torch
+Requires-Dist: torchvision
+Description-Content-Type: text/markdown
+
+# RR_layer
+
+`RR_layer` is a JAX / Equinox library implementing a **Rank-Reduction (RR) layer**: a layer that compresses its input by projecting onto a low-rank SVD basis.
+
+The mechanics are different between training and evaluation:
+
+- **During training**, the SVD basis is **recomputed from scratch for every batch**. For each batch, the layer computes the SVD of the batch, truncates to the top `r` components, and uses that truncated basis to project — so the basis is always data-dependent and up to date with the current batch, and gradients flow through the SVD/projection step. The model also automatically saves the last `N=basis_history_size` bases to be later used in evaluation.
+- **During evaluation**, when layer.eval() is called, the layer will find a common basis to the whole dataset `Uf`, based on the last `N=basis_history_size` that were saved. When the layer is called in evaluation mode, the SVD is replaced by a projection on that basis `Uf`.
+- In both cases, the shape of the input should be (N x ...), where N is the number samples in a batch. Note that the output will have the same shape is the input. This is because the layer performs the SVD, truncates, but then returns the reconstructed latent space. So the compression happens inside eventhough the output still has the same shape.
+
+## What's inside
+
+```
+RR_layer/
+├── RR_layer/
+│   ├── __init__.py
+│   └── rr_layer.py        # the layer implementation(s)
+├── examples/
+|   ├──── autoencoder
+│   |    ├── generate_data.ipynb
+│   |    ├── train_baseline.ipynb
+│   |    ├── train_RRAE.ipynb
+│   |    └── compare.ipynb
+└── tests/
+    └── tests.py
+```
+
+### `RR_layer/rr_layer.py`
+
+The library code. Contains the `RR_layer` Equinox module: a layer with two modes of operation —
+
+- **train mode**: computes the SVD of the current batch's activations, truncates to rank `r`, and projects onto that basis (gradients flow through this).
+- **eval mode**: projects onto a basis that was saved from training, rather than computing a new SVD.
+
+### `Examples/`
+
+The example notebooks are meant to be worked through **in order**:
+
+1. **`generate_data.ipynb`** — generates the dataset used in the other notebooks. Run this first; it should leave behind whatever data artifacts the later notebooks load.
+
+2. **`train_baseline.ipynb`** — trains a standard, non-RR autoencoder as a reference point. This gives you a baseline reconstruction quality and runtime to compare the RRAE against.
+
+3. **`train_RRAE.ipynb`** — trains the Rank-Reduction Auto-Encoder, built with one or more `RR_layer` in the latent space. During training this recomputes the SVD basis per batch as described above; at the end of training it should save the basis to be reused at evaluation time.
+
+4. **`compare.ipynb`** — loads the trained baseline and RRAE models (using the saved, frozen SVD basis for the RRAE) and compares them for tasks such as linear interpolation and random generation.
+
+Run the notebooks in the order above — each later notebook generally depends on artifacts (data, images) produced by the earlier ones.
+
+### `tests/tests.py`
+
+Unit tests for the `RR_layer` module. Useful both to verify correctness after installing, and as additional usage examples of the layer API beyond the notebooks (e.g. shape checks, gradient flow through the SVD step in train mode, etc.)
+
+## Running the tests
+
+```bash
+python -m pytest tests/tests.py
+```
+
+## Installation
+
+```bash
+pip install RR_layer
+```
+
+## Quick start
+
+```python
+import jax
+from RR_layer import RRLayer  # replace with actual export
+
+inp = torch.randn(32, 768)
+rr = RRLayer(rank=8)
+
+rr.train()
+y = rr(inp) # will compute the SVD
+
+rr.eval()
+y = rr(inp) # will project onto a saved basis
+```
+
+## Background
+
+For more information about Rank Reduction layers, please refer to:
+
+- **Paper:** [Rank Reduction Autoencoders](https://arxiv.org/abs/2405.13980)
+- **Related repo:** [JadM133/RRAEs](https://github.com/JadM133/RRAEs)

rr_layer-0.1.1.dist-info/RECORD

@@ -0,0 +1,7 @@
+RR_layer/__init__.py,sha256=OjDRIEzUoj_Jgnm-nZofAoJAZrioLMhuf16kYHWmgc8,31
+RR_layer/rr_layer.py,sha256=QEboZdhPMDHZvULlRLJQnO2nHD4H1h5zbpfZcSxdNk4,10027
+RR_layer/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+rr_layer-0.1.1.dist-info/METADATA,sha256=HJNXocCuUI5vbPAZSj2IJI-WibvTzWLeZn1HDiAc68Q,4369
+rr_layer-0.1.1.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+rr_layer-0.1.1.dist-info/licenses/LICENSE,sha256=QQKXj7kEw2yUJxPDzHUnEPRO7YweIU6AAlwPLnhtinA,1090
+rr_layer-0.1.1.dist-info/RECORD,,

rr_layer-0.1.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any

rr_layer-0.1.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Jad Mounayer
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.