torch-ttt 0.0.1__py3-none-any.whl

torch_ttt/engine/__init__.py

@@ -0,0 +1,5 @@
+# ruff: noqa: F401
+# from .base_engine import BaseEngine
+# from .ttt_engine import TTTEngine
+# from .ttt_pp_engine import TTTPPEngine
+# from .masked_ttt_engine import MaskedTTTEngine

torch_ttt/engine/actmad_engine.py

@@ -0,0 +1,174 @@
+import torch
+from typing import List, Dict, Any, Tuple, Union
+from contextlib import contextmanager
+
+from torch.utils.data import DataLoader
+from torch_ttt.engine.base_engine import BaseEngine
+from torch_ttt.engine_registry import EngineRegistry
+
+__all__ = ["ActMADEngine"]
+
+
+# TODO: add cuda support
+@EngineRegistry.register("actmad")
+class ActMADEngine(BaseEngine):
+    """**ActMAD** approach: multi-level pixel-wise feature alignment.
+    
+    Args:
+        model (torch.nn.Module): Model to be trained with TTT.
+        features_layer_names (List[str] | str): List of layer names to be used for feature alignment.
+        optimization_parameters (dict): The optimization parameters for the engine.
+
+    :Example:
+
+    .. code-block:: python
+
+        from torch_ttt.engine.actmad_engine import ActMADEngine
+
+        model = MyModel()
+        engine = ActMADEngine(model, ["fc1", "fc2"])
+        optimizer = torch.optim.Adam(engine.parameters(), lr=1e-4)
+
+        # Training
+        engine.train()
+        for inputs, labels in train_loader:
+            optimizer.zero_grad()
+            outputs, loss_ttt = engine(inputs)
+            loss = criterion(outputs, labels) + alpha * loss_ttt
+            loss.backward()
+            optimizer.step()
+
+        # Compute statistics for features alignment
+        engine.compute_statistics(train_loader)
+
+        # Inference
+        engine.eval()
+        for inputs, labels in test_loader:
+            output, loss_ttt = engine(inputs)
+
+    Reference:
+
+        "ActMAD: Activation Matching to Align Distributions for Test-Time Training", M. Jehanzeb Mirza, Pol Jane Soneira, Wei Lin, Mateusz Kozinski, Horst Possegger, Horst Bischof
+
+        Paper link: `PDF <https://proceedings.neurips.cc/paper/2021/hash/b618c3210e934362ac261db280128c22-Abstract.html>`_
+    """
+
+    def __init__(
+            self,
+            model: torch.nn.Module,
+            features_layer_names: Union[List[str], str],
+            optimization_parameters: Dict[str, Any] = {},
+    ):
+        super().__init__()
+        self.model = model
+        self.features_layer_names = features_layer_names
+        self.optimization_parameters = optimization_parameters
+
+        if isinstance(features_layer_names, str):
+            self.features_layer_names = [features_layer_names]
+
+        # TODO: rewrite this
+        self.target_modules = []
+        for layer_name in self.features_layer_names:
+            layer_exists = False
+            for name, module in model.named_modules():
+                if name == layer_name:
+                    layer_exists = True
+                    self.target_modules.append(module)
+                    break
+            if not layer_exists:
+                raise ValueError(f"Layer {layer_name} does not exist in the model.")
+            
+        self.reference_mean = None
+        self.reference_var = None
+
+    def ttt_forward(self, inputs) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass of the model.
+
+        Args:
+            inputs (torch.Tensor): Input tensor.
+
+        Returns:
+            Returns the current model prediction and rotation loss value.
+        """
+        with self.__capture_hook() as features_hooks:
+            outputs = self.model(inputs)
+            features = [hook.output for hook in features_hooks]
+
+        # don't compute loss during training
+        if self.training:
+            return outputs, 0
+        
+        if self.reference_var is None or self.reference_mean is None:
+            raise ValueError(
+                "Reference statistics are not computed. Please call `compute_statistics` method."
+            )
+
+        l1_loss = torch.nn.L1Loss(reduction='mean')
+        features_means = [torch.mean(feature, dim=0) for feature in features]
+        features_vars = [torch.var(feature, dim=0) for feature in features]
+
+        loss = 0
+        for i in range(len(self.target_modules)):
+            print(features_means[i].device, self.reference_mean[i].device)
+            loss += l1_loss(features_means[i], self.reference_mean[i])
+            loss += l1_loss(features_vars[i], self.reference_var[i])
+
+        return outputs, loss
+    
+    def compute_statistics(self, dataloader: DataLoader) -> None:
+        """Extract and compute reference statistics for features.
+
+        Args:
+            dataloader (DataLoader): The dataloader used for extracting features. It can return tuples of tensors, with the first element expected to be the input tensor.
+
+        Raises:
+            ValueError: If the dataloader is empty or features have mismatched dimensions.
+        """
+        
+        self.model.eval()
+        feat_stack = [[] for _ in self.target_modules]
+
+        # TODO: compute variance in more memory efficient way
+        with torch.no_grad():
+            device = next(self.model.parameters()).device
+            for sample in dataloader:
+                if len(sample) < 1:
+                    raise ValueError("Dataloader returned an empty batch.")
+                
+                inputs = sample[0].to(device)
+                with self.__capture_hook() as features_hooks:
+                    _ = self.model(inputs)
+                    features = [hook.output.cpu() for hook in features_hooks]
+
+                for i, feature in enumerate(features):
+                    feat_stack[i].append(feature)
+
+        # Compute mean and variance 
+        self.reference_mean = [torch.mean(torch.cat(feat), dim=0).to(device) for feat in feat_stack]
+        self.reference_var = [torch.var(torch.cat(feat), dim=0).to(device) for feat in feat_stack]
+
+    @contextmanager
+    def __capture_hook(self):
+        """Context manager to capture features via a forward hook."""
+    
+        class OutputHook:
+            def __init__(self):
+                self.output = None
+
+            def hook(self, module, input, output):
+                self.output = output
+
+        hook_handels = []
+        features_hooks = []
+        for module in self.target_modules:
+            hook = OutputHook()
+            features_hooks.append(hook)
+            hook_handle = module.register_forward_hook(hook.hook)
+            hook_handels.append(hook_handle)
+
+        try:
+            yield features_hooks
+        finally:
+            for hook in hook_handels:
+                hook.remove()

torch_ttt/engine/base_engine.py

@@ -0,0 +1,47 @@
+from abc import ABC, abstractmethod
+from typing import Tuple
+import torch.nn as nn
+import torch
+from copy import deepcopy
+
+class BaseEngine(nn.Module, ABC):
+
+    def __init__(self):
+        nn.Module.__init__(self)
+        self.optimization_parameters = {}
+
+    @abstractmethod
+    def ttt_forward(self, inputs) -> Tuple[torch.Tensor, torch.Tensor]:
+        pass
+
+    def forward(self, inputs):
+
+        if self.training:
+            return self.ttt_forward(inputs)
+
+        # TODO: optimization pipeline should be more flexible and 
+        # user-defined, need some special structure for that
+        optimization_parameters = self.optimization_parameters or {} 
+
+        optimizer_name = optimization_parameters.get("optimizer_name", "adam")
+        num_steps = optimization_parameters.get("num_steps", 3)
+        lr = optimization_parameters.get("lr", 1e-2)
+        copy_model = optimization_parameters.get("copy_model", "True")
+
+        running_engine = deepcopy(self) if copy_model else self
+
+        parameters = filter(lambda p: p.requires_grad, running_engine.model.parameters())
+        if optimizer_name == "adam":
+            optimizer = torch.optim.Adam(parameters, lr=lr)
+
+        loss = 0 # default value
+        for _ in range(num_steps):
+            optimizer.zero_grad()
+            _, loss = running_engine.ttt_forward(inputs)
+            loss.backward()
+            optimizer.step()
+
+        with torch.no_grad():
+            final_outputs = running_engine.model(inputs)
+
+        return final_outputs, loss

torch_ttt/engine/masked_ttt_engine.py

@@ -0,0 +1,16 @@
+from torch_ttt.engine.base_engine import BaseEngine
+from torch_ttt.engine_registry import EngineRegistry
+
+__all__ = ["MaskedTTTEngine"]
+
+
+# TODO: add cuda support
+@EngineRegistry.register("masked_ttt")
+class MaskedTTTEngine(BaseEngine):
+    """Masked autoencoders-based **test-time training** approach."""
+
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, inputs):
+        pass

torch_ttt/engine/tent_engine.py

@@ -0,0 +1,49 @@
+import torch
+from typing import Dict, Any, Tuple
+from torch_ttt.engine.base_engine import BaseEngine
+from torch_ttt.engine_registry import EngineRegistry
+
+__all__ = ["TentEngine"]
+
+@EngineRegistry.register("tent")
+class TentEngine(BaseEngine):
+    """**TENT**: Fully test-time adaptation by entropy minimization.
+
+    Args:
+        model (torch.nn.Module): The model to adapt.
+        optimization_parameters (dict): Hyperparameters for adaptation.
+
+    Reference:
+        "TENT: Fully Test-Time Adaptation by Entropy Minimization"
+        Dequan Wang, Evan Shelhamer, et al.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        optimization_parameters: Dict[str, Any] = {},
+    ):
+        super().__init__()
+        self.model = model
+        self.optimization_parameters = optimization_parameters
+
+        # Tent adapts only affine parameters in BatchNorm
+        self.model.train()
+        self._configure_bn()
+
+    def _configure_bn(self):
+        for module in self.model.modules():
+            if isinstance(module, torch.nn.BatchNorm2d):
+                module.requires_grad_(True)
+                module.track_running_stats = False
+            else:
+                for param in module.parameters(recurse=False):
+                    param.requires_grad = False
+
+    def ttt_forward(self, inputs) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass and entropy loss computation."""
+        outputs = self.model(inputs)
+        probs = torch.nn.functional.softmax(outputs, dim=1)
+        log_probs = torch.nn.functional.log_softmax(outputs, dim=1)
+        entropy = -torch.sum(probs * log_probs, dim=1).mean()
+        return outputs, entropy

torch_ttt/engine/ttt_engine.py

@@ -0,0 +1,184 @@
+import torch
+from contextlib import contextmanager
+from typing import Tuple, Dict, Any
+from torchvision.transforms import functional as F
+from torch_ttt.engine.base_engine import BaseEngine
+from torch_ttt.engine_registry import EngineRegistry
+
+__all__ = ["TTTEngine"]
+
+
+@EngineRegistry.register("ttt")
+class TTTEngine(BaseEngine):
+    r"""Original image rotation-based **test-time training** approach.
+
+    Args:
+        model (torch.nn.Module): Model to be trained with TTT.
+        features_layer_name (str): The name of the layer from which the features are extracted.
+        angle_head (torch.nn.Module, optional): The head that predicts the rotation angles.
+        angle_criterion (torch.nn.Module, optional): The loss function for the rotation angles.
+        optimization_parameters (dict): The optimization parameters for the engine.
+
+    Warning:
+        The module with the name :attr:`features_layer_name` should be present in the model.
+
+    Note:
+        :attr:`angle_head` and :attr:`angle_criterion` are optional arguments and can be user-defined. If not provided, the default shallow head and the :meth:`torch.nn.CrossEntropyLoss()` loss function are used.
+
+    Note:
+        The original `TTT <https://github.com/yueatsprograms/ttt_cifar_release/blob/acac817fb7615850d19a8f8e79930240c9afe8b5/utils/rotation.py#L27>`_ implementation uses a four-class classification task, corresponding to image rotations of 0°, 90°, 180°, and 270°.
+
+    :Example:
+
+    .. code-block:: python
+
+        from torch_ttt.engine.ttt_engine import TTTEngine
+
+        model = MyModel()
+        engine = TTTEngine(model, "fc1")
+        optimizer = torch.optim.Adam(engine.parameters(), lr=1e-4)
+
+        # Training
+        engine.train()
+        for inputs, labels in train_loader:
+            optimizer.zero_grad()
+            outputs, loss_ttt = engine(inputs)
+            loss = criterion(outputs, labels) + alpha * loss_ttt
+            loss.backward()
+            optimizer.step()
+
+        # Inference
+        engine.eval()
+        for inputs, labels in test_loader:
+            output, loss_ttt = engine(inputs)
+
+    Reference:
+
+        "Test-Time Training with Self-Supervision for Generalization under Distribution Shifts", Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei A. Efros, Moritz Hardt
+
+        Paper link: `PDF <http://proceedings.mlr.press/v119/sun20b/sun20b.pdf>`_
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        features_layer_name: str,
+        angle_head: torch.nn.Module = None,
+        angle_criterion: torch.nn.Module = None,
+        optimization_parameters: Dict[str, Any] = {},
+    ) -> None:
+        super().__init__()
+        self.model = model
+        self.angle_head = angle_head
+        self.angle_criterion = angle_criterion if angle_criterion else torch.nn.CrossEntropyLoss()
+        self.features_layer_name = features_layer_name
+        self.optimization_parameters = optimization_parameters
+
+        # Locate and store the reference to the target module
+        self.target_module = None
+        for name, module in model.named_modules():
+            if name == features_layer_name:
+                self.target_module = module
+                break
+
+        if self.target_module is None:
+            raise ValueError(f"Module '{features_layer_name}' not found in the model.")
+
+    def ttt_forward(self, inputs) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass of the model.
+
+        Args:
+            inputs (torch.Tensor): Input tensor.
+
+        Returns:
+            Returns the current model prediction and rotation loss value.
+        """
+
+        # has to dynamically register a hook to get the features and then remove it
+        # need this for deepcopying the engine, see https://github.com/pytorch/pytorch/pull/103001
+        with self.__capture_hook() as features_hook:
+            # Original forward pass, intact
+            outputs = self.model(inputs)
+
+            # See original code: https://github.com/yueatsprograms/ttt_cifar_release/blob/acac817fb7615850d19a8f8e79930240c9afe8b5/main.py#L69
+            rotated_inputs, rotation_labels = self.__rotate_inputs(inputs)
+            _ = self.model(rotated_inputs)
+            features = features_hook.output
+
+        # Build angle head if not already built
+        if self.angle_head is None:
+            self.angle_head = self.__build_angle_head(features)
+
+        # move angle head to the same device as the features
+        self.angle_head.to(features.device)
+        angles = self.angle_head(features)
+
+        # Compute rotation loss
+        rotation_loss = self.angle_criterion(angles, rotation_labels)
+        return outputs, rotation_loss
+
+    # Follow this code (expand case): https://github.com/yueatsprograms/ttt_cifar_release/blob/acac817fb7615850d19a8f8e79930240c9afe8b5/utils/rotation.py#L27
+    def __rotate_inputs(self, inputs) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Rotate the input images by 0, 90, 180, and 270 degrees."""
+        device = next(self.model.parameters()).device
+        rotated_image_90 = F.rotate(inputs, 90)
+        rotated_image_180 = F.rotate(inputs, 180)
+        rotated_image_270 = F.rotate(inputs, 270)
+        batch_size = inputs.shape[0]
+        inputs = torch.cat([inputs, rotated_image_90, rotated_image_180, rotated_image_270], dim=0)
+        labels = [0] * batch_size + [1] * batch_size + [2] * batch_size + [3] * batch_size
+        return inputs.to(device), torch.tensor(labels, dtype=torch.long).to(device)
+
+    def __build_angle_head(self, features) -> torch.nn.Module:
+        """Build the angle head."""
+        device = next(self.model.parameters()).device
+
+        # See original implementation: https://github.com/yueatsprograms/ttt_cifar_release/blob/acac817fb7615850d19a8f8e79930240c9afe8b5/utils/test_helpers.py#L33C10-L33C39
+        if len(features.shape) == 2:
+            return torch.nn.Sequential(
+                torch.nn.Linear(features.shape[1], 16),
+                torch.nn.ReLU(),
+                torch.nn.Linear(16, 8),
+                torch.nn.ReLU(),
+                torch.nn.Linear(8, 4),
+            ).to(device)
+
+        # See original implementation: https://github.com/yueatsprograms/ttt_cifar_release/blob/acac817fb7615850d19a8f8e79930240c9afe8b5/models/SSHead.py#L29
+        elif len(features.shape) == 4:
+            return torch.nn.Sequential(
+                torch.nn.Conv2d(features.shape[1], 16, 3),
+                torch.nn.ReLU(),
+                torch.nn.Conv2d(16, 4, 3),
+                torch.nn.AdaptiveAvgPool2d((1, 1)),  # Global Average Pooling
+                torch.nn.Flatten(),
+            ).to(device)
+
+        elif len(features.shape) == 5:  # For 3D inputs (batch, channels, depth, height, width)
+            return torch.nn.Sequential(
+                torch.nn.Conv3d(features.shape[1], 16, kernel_size=3),
+                torch.nn.ReLU(),
+                torch.nn.Conv3d(16, 4, kernel_size=3),
+                torch.nn.AdaptiveAvgPool3d((1, 1, 1)),  # Global Average Pooling
+                torch.nn.Flatten(),
+            ).to(device)
+
+        raise ValueError("Invalid input tensor shape.")
+
+    @contextmanager
+    def __capture_hook(self):
+        """Context manager to capture features via a forward hook."""
+
+        class OutputHook:
+            def __init__(self):
+                self.output = None
+
+            def hook(self, module, input, output):
+                self.output = output
+
+        features_hook = OutputHook()
+        hook_handle = self.target_module.register_forward_hook(features_hook.hook)
+
+        try:
+            yield features_hook
+        finally:
+            hook_handle.remove()

torch_ttt/engine/ttt_pp_engine.py

@@ -0,0 +1,291 @@
+import torch
+from typing import Tuple, Optional, Callable, Dict, Any
+from contextlib import contextmanager
+
+from torchvision import transforms
+from torch.utils.data import DataLoader
+from torch_ttt.engine.base_engine import BaseEngine
+from torch_ttt.engine_registry import EngineRegistry
+from torch_ttt.loss.contrastive_loss import ContrastiveLoss
+from torch_ttt.utils.augmentations import RandomResizedCrop
+
+__all__ = ["TTTPPEngine"]
+
+
+# TODO: finish this class
+@EngineRegistry.register("ttt_pp")
+class TTTPPEngine(BaseEngine):
+    """**TTT++** approach: feature alignment-based + SimCLR loss.
+
+    Args:
+        model (torch.nn.Module): Model to be trained with TTT.
+        features_layer_name (str): The name of the layer from which the features are extracted.
+        contrastive_head (torch.nn.Module, optional): The head that is used for SimCLR's Loss.
+        contrastive_criterion (torch.nn.Module, optional): The loss function used for SimCLR.
+        contrastive_transform (callable): A transformation or a composition of transformations applied to the input images to generate augmented views for contrastive learning.
+        scale_cov (float): The scale factor for the covariance loss.
+        scale_mu (float): The scale factor for the mean loss.
+        scale_c_cov (float): The scale factor for the contrastive covariance loss.
+        scale_c_mu (float): The scale factor for the contrastive mean loss.
+        optimization_parameters (dict): The optimization parameters for the engine.
+    Warning:
+        The module with the name :attr:`features_layer_name` should be present in the model.
+
+    :Example:
+
+    .. code-block:: python
+
+        from torch_ttt.engine.ttt_pp_engine import TTTPPEngine
+
+        model = MyModel()
+        engine = TTTPPEngine(model, "fc1")
+        optimizer = torch.optim.Adam(engine.parameters(), lr=1e-4)
+
+        # Training
+        engine.train()
+        for inputs, labels in train_loader:
+            optimizer.zero_grad()
+            outputs, loss_ttt = engine(inputs)
+            loss = criterion(outputs, labels) + alpha * loss_ttt
+            loss.backward()
+            optimizer.step()
+
+        # Compute statistics for features alignment
+        engine.compute_statistics(train_loader)
+
+        # Inference
+        engine.eval()
+        for inputs, labels in test_loader:
+            output, loss_ttt = engine(inputs)
+
+    Reference:
+
+        "TTT++: When Does Self-Supervised Test-Time Training Fail or Thrive?", Yuejiang Liu, Parth Kothari, Bastien van Delft, Baptiste Bellot-Gurlet, Taylor Mordan, Alexandre Alahi
+
+        Paper link: `PDF <https://proceedings.neurips.cc/paper/2021/hash/b618c3210e934362ac261db280128c22-Abstract.html>`_
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        features_layer_name: str,
+        contrastive_head: torch.nn.Module = None,
+        contrastive_criterion: torch.nn.Module = ContrastiveLoss(),
+        contrastive_transform: Optional[Callable] = None,
+        scale_cov: float = 0.1,
+        scale_mu: float = 0.1,
+        scale_c_cov: float = 0.1,
+        scale_c_mu: float = 0.1,
+        optimization_parameters: Dict[str, Any] = {},
+    ) -> None:
+        super().__init__()
+        self.model = model
+        self.features_layer_name = features_layer_name
+        self.contrastive_head = contrastive_head
+        self.contrastive_criterion = (
+            contrastive_criterion if contrastive_criterion else ContrastiveLoss()
+        )
+        self.scale_cov = scale_cov
+        self.scale_mu = scale_mu
+        self.scale_c_cov = scale_c_cov
+        self.scale_c_mu = scale_c_mu
+        self.contrastive_transform = contrastive_transform
+
+        self.reference_cov = None
+        self.reference_mean = None
+        self.reference_c_cov = None
+        self.reference_c_mean = None
+
+        self.optimization_parameters = optimization_parameters
+
+        # Locate and store the reference to the target module
+        self.target_module = None
+        for name, module in model.named_modules():
+            if name == features_layer_name:
+                self.target_module = module
+                break
+
+        if self.target_module is None:
+            raise ValueError(f"Module '{features_layer_name}' not found in the model.")
+
+        # Validate that the target module is a Linear layer
+        if not isinstance(self.target_module, torch.nn.Linear):
+            raise TypeError(
+                f"Module '{features_layer_name}' is expected to be of type 'torch.nn.Linear', "
+                f"but found type '{type(self.target_module).__name__}'."
+            )
+
+        if contrastive_transform is None:
+            # default SimCLR augmentation
+            self.contrastive_transform = transforms.Compose(
+                [
+                    RandomResizedCrop(scale=(0.2, 1.0)),
+                    transforms.RandomGrayscale(p=0.2),
+                    transforms.RandomApply([transforms.GaussianBlur(5)], p=0.3),
+                    transforms.RandomHorizontalFlip(),
+                ]
+            )
+
+    def __build_contrastive_head(self, features) -> torch.nn.Module:
+        """Build the angle head."""
+        device = next(self.model.parameters()).device
+        if len(features.shape) == 2:
+            return torch.nn.Sequential(
+                torch.nn.Linear(features.shape[1], 16),
+                torch.nn.ReLU(),
+                torch.nn.Linear(16, 16),
+                torch.nn.ReLU(),
+                torch.nn.Linear(16, 16),
+            ).to(device)
+
+        raise ValueError("Features should be 2D tensor.")
+
+    def ttt_forward(self, inputs) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass of the model.
+
+        Args:
+            inputs (torch.Tensor): Input tensor.
+
+        Returns:
+            Returns the model prediction and TTT++ loss value.
+        """
+
+        # reset reference statistics during training
+        if self.training:
+            self.reference_cov = None
+            self.reference_mean = None
+            self.reference_c_cov = None
+            self.reference_c_mean = None
+
+        contrastive_inputs = torch.cat(
+            [self.contrastive_transform(inputs), self.contrastive_transform(inputs)], dim=0
+        )
+
+        # extract features for contrastive loss
+        with self.__capture_hook() as features_hook:
+            _ = self.model(contrastive_inputs)
+            features = features_hook.output
+
+        # Build angle head if not already built
+        if self.contrastive_head is None:
+            self.contrastive_head = self.__build_contrastive_head(features)
+
+        contrasitve_features = self.contrastive_head(features)
+        contrasitve_features = contrasitve_features.view(2, len(inputs), -1).transpose(0, 1)
+        loss = self.contrastive_criterion(contrasitve_features)
+
+        # make inference for a final prediction
+        with self.__capture_hook() as features_hook:
+            outputs = self.model(inputs)
+            features = features_hook.output
+
+        # compute alignment loss only during test
+        if not self.training:
+            if (
+                self.reference_cov is None
+                or self.reference_mean is None
+                or self.reference_c_cov is None
+                or self.reference_c_mean is None
+            ):
+                raise ValueError(
+                    "Reference statistics are not computed. Please call `compute_statistics` method."
+                )
+
+            # compute features alignment loss
+            cov_ext = self.__covariance(features)
+            mu_ext = features.mean(dim=0)
+
+            d = self.reference_cov.shape[0]
+
+            loss += self.scale_cov * (self.reference_cov - cov_ext).pow(2).sum() / (4.0 * d**2)
+            loss += self.scale_mu * (self.reference_mean - mu_ext).pow(2).mean()
+
+            # compute contrastive features alignment loss
+            c_features = self.contrastive_head(features)
+
+            cov_ext = self.__covariance(c_features)
+            mu_ext = c_features.mean(dim=0)
+
+            d = self.reference_c_cov.shape[0]
+            loss += self.scale_c_cov * (self.reference_c_cov - cov_ext).pow(2).sum() / (4.0 * d**2)
+            loss += self.scale_c_mu * (self.reference_c_mean - mu_ext).pow(2).mean()
+
+        return outputs, loss
+
+    @staticmethod
+    def __covariance(features):
+        """Legacy wrapper to maintain compatibility in the engine."""
+        from torch_ttt.utils.math import compute_covariance
+
+        return compute_covariance(features, dim=0)
+
+    def compute_statistics(self, dataloader: DataLoader) -> None:
+        """Extract and compute reference statistics for features and contrastive features.
+
+        Args:
+            dataloader (DataLoader): The dataloader used for extracting features. It can return tuples of tensors, with the first element expected to be the input tensor.
+
+        Raises:
+            ValueError: If the dataloader is empty or features have mismatched dimensions.
+        """
+
+        self.model.eval()
+
+        feat_stack = []
+        c_feat_stack = []
+
+        with torch.no_grad():
+            device = next(self.model.parameters()).device
+            for sample in dataloader:
+                if len(sample) < 1:
+                    raise ValueError("Dataloader returned an empty batch.")
+
+                inputs = sample[0].to(device)
+                with self.__capture_hook() as features_hook:
+                    _ = self.model(inputs)
+                    feat = features_hook.output
+
+                    # Initialize contrastive head if not already initialized
+                    if self.contrastive_head is None:
+                        self.contrastive_head = self.__build_contrastive_head(feat)
+
+                    # Compute contrastive features
+                    contrastive_feat = self.contrastive_head(feat)
+
+                feat_stack.append(feat.cpu())
+                c_feat_stack.append(contrastive_feat.cpu())
+
+        # compute features statistics
+        feat_all = torch.cat(feat_stack)
+        feat_cov = self.__covariance(feat_all)
+        feat_mean = feat_all.mean(dim=0)
+
+        self.reference_cov = feat_cov.to(device)
+        self.reference_mean = feat_mean.to(device)
+
+        # compute contrastive features statistics
+        feat_all = torch.cat(c_feat_stack)
+        feat_cov = self.__covariance(feat_all)
+        feat_mean = feat_all.mean(dim=0)
+
+        self.reference_c_cov = feat_cov.to(device)
+        self.reference_c_mean = feat_mean.to(device)
+
+    @contextmanager
+    def __capture_hook(self):
+        """Context manager to capture features via a forward hook."""
+
+        class OutputHook:
+            def __init__(self):
+                self.output = None
+
+            def hook(self, module, input, output):
+                self.output = output
+
+        features_hook = OutputHook()
+        hook_handle = self.target_module.register_forward_hook(features_hook.hook)
+
+        try:
+            yield features_hook
+        finally:
+            hook_handle.remove()

torch_ttt/engine_registry.py

@@ -0,0 +1,38 @@
+import importlib
+import os
+from torch_ttt.engine.base_engine import BaseEngine
+
+
+class EngineRegistry:
+    _registry = {}
+
+    @classmethod
+    def register(cls, name):
+        def decorator(engine_class):
+            if not issubclass(engine_class, BaseEngine):
+                raise TypeError(f"Engine class '{name}' must inherit from BaseEngine.")
+            if name in cls._registry:
+                raise ValueError(f"Engine '{name}' is already registered.")
+
+            cls._registry[name] = engine_class
+            return engine_class
+
+        return decorator
+
+    @classmethod
+    def get_engine(cls, name):
+        if name not in cls._registry:
+            raise ValueError(f"Engine '{name}' is not registered.")
+        return cls._registry[name]
+
+
+# Dynamically import all losses in the losses directory
+def register_all_engines():
+    losses_dir = os.path.dirname(__file__) + "/engine"
+    for file in os.listdir(losses_dir):
+        if file.endswith("_engine.py") and not file.startswith("__"):
+            module_name = f"torch_ttt.engine.{file[:-3]}"
+            importlib.import_module(module_name)
+
+
+register_all_engines()

torch_ttt/loss/__init__.py

@@ -0,0 +1,2 @@
+# ruff: noqa: F401
+from .base_loss import BaseLoss

torch_ttt/loss/base_loss.py

@@ -0,0 +1,8 @@
+from abc import ABC, abstractmethod
+import torch.nn as nn
+
+
+class BaseLoss(nn.Module, ABC):
+    @abstractmethod
+    def forward(self, model, inputs):
+        pass

torch_ttt/loss/contrastive_loss.py

@@ -0,0 +1,94 @@
+import torch
+from torch_ttt.loss.base_loss import BaseLoss
+from torch_ttt.loss_registry import LossRegistry
+
+
+# TODO: make it more readable
+@LossRegistry.register("contrastive")
+class ContrastiveLoss(BaseLoss):
+    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
+    It also supports the unsupervised contrastive loss in SimCLR"""
+
+    def __init__(self, temperature=0.07, contrast_mode="all", base_temperature=0.07) -> None:
+        super(ContrastiveLoss, self).__init__()
+        self.temperature = temperature
+        self.contrast_mode = contrast_mode
+        self.base_temperature = base_temperature
+
+    def forward(self, features, labels=None, mask=None):
+        """Compute loss for model. If both `labels` and `mask` are None,
+        it degenerates to SimCLR unsupervised loss:
+        https://arxiv.org/pdf/2002.05709.pdf
+
+        Args:
+            features: hidden vector of shape [bsz, n_views, ...].
+            labels: ground truth of shape [bsz].
+            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
+                has the same class as sample i. Can be asymmetric.
+        Returns:
+            A loss scalar.
+        """
+        device = torch.device("cuda") if features.is_cuda else torch.device("cpu")
+
+        if len(features.shape) < 3:
+            raise ValueError(
+                "`features` needs to be [bsz, n_views, ...]," "at least 3 dimensions are required"
+            )
+        if len(features.shape) > 3:
+            features = features.view(features.shape[0], features.shape[1], -1)
+
+        batch_size = features.shape[0]
+        if labels is not None and mask is not None:
+            raise ValueError("Cannot define both `labels` and `mask`")
+        elif labels is None and mask is None:
+            mask = torch.eye(batch_size, dtype=torch.float32).to(device)
+        elif labels is not None:
+            labels = labels.contiguous().view(-1, 1)
+            if labels.shape[0] != batch_size:
+                raise ValueError("Num of labels does not match num of features")
+            mask = torch.eq(labels, labels.T).float().to(device)
+        else:
+            mask = mask.float().to(device)
+
+        contrast_count = features.shape[1]
+        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
+        if self.contrast_mode == "one":
+            anchor_feature = features[:, 0]
+            anchor_count = 1
+        elif self.contrast_mode == "all":
+            anchor_feature = contrast_feature
+            anchor_count = contrast_count
+        else:
+            raise ValueError("Unknown mode: {}".format(self.contrast_mode))
+
+        # compute logits
+        anchor_dot_contrast = torch.div(
+            torch.matmul(anchor_feature, contrast_feature.T), self.temperature
+        )
+        # for numerical stability
+        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
+        logits = anchor_dot_contrast - logits_max.detach()
+
+        # tile mask
+        mask = mask.repeat(anchor_count, contrast_count)
+        # mask-out self-contrast cases
+        logits_mask = torch.scatter(
+            torch.ones_like(mask),
+            1,
+            torch.arange(batch_size * anchor_count).view(-1, 1).to(device),
+            0,
+        )
+        mask = mask * logits_mask
+
+        # compute log_prob
+        exp_logits = torch.exp(logits) * logits_mask
+        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
+
+        # compute mean of log-likelihood over positive
+        mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)
+
+        # loss
+        loss = -(self.temperature / self.base_temperature) * mean_log_prob_pos
+        loss = loss.view(anchor_count, batch_size).mean()
+
+        return loss

torch_ttt/loss/entropy_loss.py

@@ -0,0 +1,15 @@
+import torch
+import torch.nn.functional as F
+from torch_ttt.loss.base_loss import BaseLoss
+from torch_ttt.loss_registry import LossRegistry
+
+
+@LossRegistry.register("entropy")
+class EntropyLoss(BaseLoss):
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, model, inputs):
+        logits = model(inputs)
+        probs = F.softmax(logits, dim=1)
+        return -torch.sum(probs * torch.log(probs), dim=1).mean()

torch_ttt/loss/mean_loss.py

@@ -0,0 +1,12 @@
+from torch_ttt.loss.base_loss import BaseLoss
+from torch_ttt.loss_registry import LossRegistry
+
+
+@LossRegistry.register("mean")
+class MeanLoss(BaseLoss):
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, model, inputs):
+        outputs = model(inputs)
+        return outputs.mean()

torch_ttt/loss/ttt_loss.py

@@ -0,0 +1,13 @@
+from torch_ttt.loss.base_loss import BaseLoss
+from torch_ttt.loss_registry import LossRegistry
+
+
+@LossRegistry.register("ttt")
+class TTTLoss(BaseLoss):
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, model, inputs):
+        # TODO [P2]: check that model is TTTEngine
+        _, loss = model(inputs)
+        return loss

torch_ttt/loss/weights_magnitude_loss.py

@@ -0,0 +1,30 @@
+import torch
+from torch_ttt.loss.base_loss import BaseLoss
+from torch_ttt.loss_registry import LossRegistry
+
+
+@LossRegistry.register("weights_magnitude")
+class WeightsMagnitudeLoss(BaseLoss):
+    def __init__(self):
+        super().__init__()
+        self.quantile = 0.95
+
+    def __call__(self, model, inputs):
+        # Step 2: Collect all model weights
+        all_weights = []
+        for param in model.parameters():
+            if param.requires_grad:  # Focus only on trainable parameters
+                all_weights.append(param.view(-1))  # Flatten weights into a 1D tensor
+
+        # Concatenate all weights into a single tensor
+        all_weights = torch.cat(all_weights)
+
+        # Step 3: Compute the top `quantile` values
+        quantile_value = torch.quantile(all_weights.abs(), self.quantile)
+        top_quantile_weights = all_weights[all_weights.abs() >= quantile_value]
+
+        # Step 4: Compute the average of top quantile weights
+        weight_loss = top_quantile_weights.mean()
+
+        # Return the combined loss
+        return weight_loss

torch_ttt/loss/zerot_loss.py

@@ -0,0 +1,158 @@
+import torch
+from torch_ttt.loss.base_loss import BaseLoss
+from torch_ttt.loss_registry import LossRegistry
+
+
+@LossRegistry.register("zerot")
+class ZeroTrainLoss(BaseLoss):
+    def __init__(self):
+        super().__init__()
+        self.quantile = 0.95
+
+    def __call__(self, model, inputs):
+        N = len(list(model.named_parameters()))
+
+        importance_dict = compute_weight_importance(model, inputs, N)
+
+        # Calculate the top s% mean importance as the loss
+        loss = top_s_percent_mean(importance_dict, self.quantile)
+
+        # Return the combined loss
+        return loss
+
+
+def top_s_percent_mean(importance_dict, s):
+    """
+    Compute the mean of the top s% of the importance values.
+
+    Args:
+        importance_dict (dict): A dictionary mapping layer names to weight importance tensors.
+        s (float): The percentage (0-100) of top importance values to consider.
+
+    Returns:
+        torch.Tensor: The mean of the top s% importance values, differentiable with respect to original weights.
+    """
+    # Concatenate all importance values into a single tensor
+    all_importances = torch.cat([importance.view(-1) for importance in importance_dict.values()])
+
+    # Determine the number of top elements to keep
+    k = max(
+        1, int(s / 100.0 * all_importances.numel())
+    )  # Ensure at least one element is considered
+
+    # Sort the importance values in descending order
+    top_importances, _ = torch.topk(all_importances, k)
+
+    # Compute the mean of the top s% importance values
+    top_mean = top_importances.mean()
+
+    return top_mean
+
+
+def optimize_last_n_layers(m, x, N, s, M, lr=1e-3, optimizer=None, verbose=False):
+    """
+    Optimize the last N layers of the model for M steps using the top s% of weight importance as a TTT loss function.
+
+    Args:
+        m (torch.nn.Module): The model.
+        x (torch.Tensor): The input batch.
+        N (int): The number of layers from the end to optimize.
+        s (float): The percentage (0-100) of top importance values to consider for the loss.
+        M (int): The number of optimization steps.
+        lr (float, optional): Learning rate for the optimizer. Default is 1e-3.
+
+    Returns:
+        torch.Tensor: The model's predictions for the input batch after optimization.
+    """
+    # Set the model to training mode
+    m.train()
+
+    # Collect parameters of the last N layers to optimize
+    layers = list(m.named_parameters())[-2 * N :]
+    params_to_optimize = [param for name, param in layers if "weight" in name]
+
+    # Ensure that the parameters of the last N layers have requires_grad=True
+    for param in params_to_optimize:
+        param.requires_grad = True
+
+    # Set up the optimizer for the last N layers' parameters
+    if optimizer is None:
+        optimizer = torch.optim.Adam(params_to_optimize, lr=lr)
+
+    # Optimization loop for M steps
+    for step in range(M):
+        # Compute weight importance
+        importance_dict = compute_weight_importance(m, x, N)
+
+        # Calculate the top s% mean importance as the loss
+        loss = top_s_percent_mean(importance_dict, s)
+
+        # Zero the gradients
+        optimizer.zero_grad()
+
+        # Backward pass to compute gradients
+        loss.backward()
+
+        # Step the optimizer to update parameters
+        optimizer.step()
+
+        # Optionally, print loss for monitoring
+        if verbose:
+            print(f"Step {step+1}/{M}, Loss: {loss.item()}")
+
+    # After optimization, run the model on the input x to get predictions
+    m.eval()  # Set the model to evaluation mode
+    with torch.no_grad():  # Disable gradient computation for inference
+        predictions = m(x)
+
+    return predictions, optimizer
+
+
+def compute_weight_importance(m, x, N):
+    """
+    Compute the weight importance for the last N layers of the model efficiently.
+
+    Args:
+        m (torch.nn.Module): The model.
+        x (torch.Tensor): The input batch.
+        N (int): The number of layers from the end to consider.
+
+    Returns:
+        dict: A dictionary mapping layer names to the weight importance tensor.
+    """
+    # Ensure the model is in evaluation mode
+    m.eval()
+
+    # Disable gradients for all parameters initially
+    for param in m.parameters():
+        param.requires_grad = False
+
+    # Enable gradients only for the last N layers
+    layers = list(m.named_parameters())[-2 * N :]  # Get the last N layers (both weights and biases)
+    for name, param in layers:
+        if "weight" in name:
+            param.requires_grad = True
+
+    # Forward pass to compute the output
+    output = m(x)
+
+    # Use a simple loss function (sum of outputs) to create a scalar output
+    loss = output.mean()
+
+    # Compute gradients with respect to the last N layers' parameters
+    gradients = torch.autograd.grad(
+        loss, [param for name, param in layers if "weight" in name], create_graph=True
+    )
+
+    # Calculate importance: w * \nabla{w} for the last N layers
+    importance_dict = {}
+    grad_idx = 0
+    for name, param in layers:
+        if "weight" in name:
+            importance = (param * gradients[grad_idx]).abs()  # w * \nabla{w}
+            # importance = (param).abs()  # w * \nabla{w}
+            # importance = gradients[grad_idx].abs()  #  # w * \nabla{w}
+            importance_dict[name] = importance  # Don't detach the importance
+            grad_idx += 1
+
+    return importance_dict

torch_ttt/loss_registry.py

@@ -0,0 +1,38 @@
+import importlib
+import os
+from torch_ttt.loss.base_loss import BaseLoss
+
+
+class LossRegistry:
+    _registry = {}
+
+    @classmethod
+    def register(cls, name):
+        def decorator(loss_class):
+            if not issubclass(loss_class, BaseLoss):
+                raise TypeError(f"Loss class '{name}' must inherit from BaseLoss.")
+            if name in cls._registry:
+                raise ValueError(f"Loss '{name}' is already registered.")
+
+            cls._registry[name] = loss_class
+            return loss_class
+
+        return decorator
+
+    @classmethod
+    def get_loss(cls, name):
+        if name not in cls._registry:
+            raise ValueError(f"Loss '{name}' is not registered.")
+        return cls._registry[name]
+
+
+# Dynamically import all losses in the losses directory
+def register_all_losses():
+    losses_dir = os.path.dirname(__file__) + "/loss"
+    for file in os.listdir(losses_dir):
+        if file.endswith("_loss.py") and not file.startswith("__"):
+            module_name = f"torch_ttt.loss.{file[:-3]}"
+            importlib.import_module(module_name)
+
+
+register_all_losses()

torch_ttt/utils/augmentations.py

@@ -0,0 +1,25 @@
+import random
+from torchvision.transforms import functional as F
+
+
+class RandomResizedCrop:
+    def __init__(self, scale=(0.2, 1.0)):
+        self.scale = scale
+
+    def __call__(self, img):
+        # Dynamically compute the crop size
+        original_size = img.shape[-2:]  # H × W
+        area = original_size[0] * original_size[1]
+        target_area = random.uniform(self.scale[0], self.scale[1]) * area
+        aspect_ratio = random.uniform(3 / 4, 4 / 3)
+
+        h = int(round((target_area * aspect_ratio) ** 0.5))
+        w = int(round((target_area / aspect_ratio) ** 0.5))
+
+        if random.random() < 0.5:  # Randomly swap h and w
+            h, w = w, h
+
+        h = min(h, original_size[0])
+        w = min(w, original_size[1])
+
+        return F.resized_crop(img, top=0, left=0, height=h, width=w, size=original_size)

torch_ttt/utils/math.py

@@ -0,0 +1,29 @@
+import torch
+
+
+def compute_covariance(features: torch.Tensor, dim: int = 0) -> torch.Tensor:
+    """Compute covariance matrix for given features along a specific dimension.
+
+    Args:
+        features (torch.Tensor): Input tensor of shape [N, D] or higher dimensions.
+        dim (int): The dimension along which to compute covariance.
+
+    Returns:
+        torch.Tensor: Covariance matrix of shape [D, D].
+
+    Raises:
+        ValueError: If the input tensor has fewer than 2 dimensions.
+    """
+    if features.ndim < 2:
+        raise ValueError("Input tensor must have at least 2 dimensions to compute covariance.")
+
+    if features.size(dim) <= 1:
+        raise ValueError(
+            f"Cannot compute covariance with less than 2 samples along dimension {dim}."
+        )
+
+    n = features.shape[0]
+    tmp = torch.ones((1, n), device=features.device) @ features
+    cov = (features.t() @ features - (tmp.t() @ tmp) / n) / (n - 1)
+
+    return cov

torch_ttt-0.0.1.dist-info/METADATA

@@ -0,0 +1,16 @@
+Metadata-Version: 2.4
+Name: torch-ttt
+Version: 0.0.1
+License-File: LICENSE
+Provides-Extra: dev
+Requires-Dist: pytest-cov; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: sphinxawesome-theme==6.0.0b1; extra == "docs"
+Requires-Dist: sphinx-gallery==0.18.0; extra == "docs"
+Requires-Dist: sphinxcontrib-bibtex; extra == "docs"
+Requires-Dist: sphinxcontrib-googleanalytics; extra == "docs"
+Requires-Dist: matplotlib; extra == "docs"
+Requires-Dist: tqdm; extra == "docs"
+Provides-Extra: all
+Requires-Dist: torch_ttt[dev,docs]; extra == "all"
+Dynamic: license-file

torch_ttt-0.0.1.dist-info/RECORD

@@ -0,0 +1,26 @@
+torch_ttt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torch_ttt/engine_registry.py,sha256=FzjHxBDCF5y8vkToKDyrLcvVNAtw2LUu4Wl-EViNHw4,1144
+torch_ttt/loss_registry.py,sha256=PGkrOFzeTKzXFoDLHUf-rH8aBAtyxrh4xzyowhVQtv4,1108
+torch_ttt/engine/__init__.py,sha256=GlxEg_jHxKmyy-uLMIrZfbsPeF9hPDclV68fTIio3zU,183
+torch_ttt/engine/actmad_engine.py,sha256=tiQLBO2SzYljUzBxdHywzrN188AsfPAqV5X6CQy7dcs,6281
+torch_ttt/engine/base_engine.py,sha256=sQZq_NDzNlQnBERlRKD43Wo8Z7FJnblwHn0wANe6Q0o,1527
+torch_ttt/engine/masked_ttt_engine.py,sha256=X-SApdbzzcViOk7lijPca3IMV1wP-8QDGWqWXow5RHo,404
+torch_ttt/engine/tent_engine.py,sha256=w5JIL5eeQMO_NFlzEV92xL0MwOWExN2-8t7GdRP63Jk,1686
+torch_ttt/engine/ttt_engine.py,sha256=pzdP4vgxAz-MQwP3KJUF0xFcUpm_hf0I_YzWYx-6q60,7675
+torch_ttt/engine/ttt_pp_engine.py,sha256=W_6vP7Oc1eeVUBAix2qT3-Hz0JcV11uOdQQsfSMLJC8,11014
+torch_ttt/loss/__init__.py,sha256=rV6RxTWgU_BtdUX0d7UXhtlTexhIiqxVNm4eDO9HUOU,51
+torch_ttt/loss/base_loss.py,sha256=O7UhmrK8NGVI8Dd8zCgh_MKh5HNAkYnyq9_aAesbbpg,163
+torch_ttt/loss/contrastive_loss.py,sha256=uxlDfSqrJEmg8v8TqJgK3LhjcRsF9euGIEGG2sYS1iY,3753
+torch_ttt/loss/entropy_loss.py,sha256=bM9LetJXic5FSs3kJ6prjXkAkc6xEBaJezfCKx4-bPg,434
+torch_ttt/loss/mean_loss.py,sha256=iZZ8DvU3-2nT3D0kgXBaSA7oO74rXrp9jrS-QbLUOe0,307
+torch_ttt/loss/ttt_loss.py,sha256=B83lwl5vOZKn_iKL7OznD3BOR9bf3X4OfCz2jzKXF3Q,346
+torch_ttt/loss/weights_magnitude_loss.py,sha256=MpyW7fqvoFzwT-nQMZN_uisf1n956knzpZWd8yXmn7s,1067
+torch_ttt/loss/zerot_loss.py,sha256=k46u3RrDdF44-f_FpIyAS-r-ZsAHOU_qKp733lDn_VI,5388
+torch_ttt/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torch_ttt/utils/augmentations.py,sha256=y-thszixcgGvGFtCzMeG5pWvNVr1D_37kqyNJfd9ApE,831
+torch_ttt/utils/math.py,sha256=mmgsyCOWGXP21gqQyrHhj-axmI4GCQ1AxFCPvvrO5Uk,952
+torch_ttt-0.0.1.dist-info/licenses/LICENSE,sha256=ppPMBj3EeJB4NHrRgxWoRSoVIxx3apK0BFsIB2-jcco,1072
+torch_ttt-0.0.1.dist-info/METADATA,sha256=kWAuJPGHukr9eKsvgeDvCIZAZkIgIG2c1OiNkr-gnBQ,562
+torch_ttt-0.0.1.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+torch_ttt-0.0.1.dist-info/top_level.txt,sha256=WP1U2F1ULE-SjEdozpFfUhCR2ZNqIi8lKvUoeQ5G5Zw,10
+torch_ttt-0.0.1.dist-info/RECORD,,

torch_ttt-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (79.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

torch_ttt-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Nikita Durasov 
+
+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.

torch_ttt-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+torch_ttt