nshtrainer 1.0.0b39__tar.gz → 1.0.0b41__tar.gz

{nshtrainer-1.0.0b39 → nshtrainer-1.0.0b41}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.0.0b39
+Version: 1.0.0b41
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.0.0b39 → nshtrainer-1.0.0b41}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "1.0.0-beta39"
+version = "1.0.0-beta41"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -29,6 +29,8 @@ pyright = "*"
 ruff = "*"
 ipykernel = "*"
 ipywidgets = "*"
+pytest = "^8.3.5"
+pytest-cov = "^6.0.0"
 
 [build-system]
 requires = ["poetry-core"]
@@ -58,3 +60,19 @@ extra = [
     "huggingface-hub",
     "nshutils",
 ]
+
+[tool.pytest]
+testpaths = ["tests"]
+python_files = "test_*.py"
+python_functions = "test_*"
+python_classes = "Test*"
+
+[tool.pytest.ini_options]
+minversion = "6.0"
+# addopts = "--cov=src/nshtrainer --cov-report=term-missing --cov-report=xml --"
+addopts = [
+    "--import-mode=importlib",
+    "--cov=src/nshtrainer",
+    "--cov-report=term-missing",
+    "--cov-report=xml",
+]

nshtrainer-1.0.0b41/src/nshtrainer/configs/.gitattributes

@@ -0,0 +1 @@
+* linguist-generated=true

{nshtrainer-1.0.0b39 → nshtrainer-1.0.0b41}/src/nshtrainer/nn/__init__.py

@@ -4,6 +4,7 @@ from .mlp import MLP as MLP
 from .mlp import MLPConfig as MLPConfig
 from .mlp import MLPConfigDict as MLPConfigDict
 from .mlp import ResidualSequential as ResidualSequential
+from .mlp import custom_seed_context as custom_seed_context
 from .module_dict import TypedModuleDict as TypedModuleDict
 from .module_list import TypedModuleList as TypedModuleList
 from .nonlinearity import ELUNonlinearityConfig as ELUNonlinearityConfig

{nshtrainer-1.0.0b39 → nshtrainer-1.0.0b41}/src/nshtrainer/nn/mlp.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 import copy
 from collections.abc import Callable, Sequence
 from typing import Literal, Protocol, runtime_checkable
@@ -44,6 +45,9 @@ class MLPConfigDict(TypedDict):
     residual: bool
     """Whether to use residual connections between layers."""
 
+    seed: int | None
+    """Random seed to use for initialization. If None, the default Torch behavior is used."""
+
 
 class MLPConfig(C.Config):
     bias: bool = True
@@ -64,15 +68,20 @@ class MLPConfig(C.Config):
     residual: bool = False
     """Whether to use residual connections between layers."""
 
+    seed: int | None = None
+    """Random seed to use for initialization. If None, the default Torch behavior is used."""
+
     def to_kwargs(self) -> MLPConfigDict:
-        return {
+        kwargs: MLPConfigDict = {
             "bias": self.bias,
             "no_bias_scalar": self.no_bias_scalar,
             "nonlinearity": self.nonlinearity,
             "ln": self.ln,
             "dropout": self.dropout,
             "residual": self.residual,
+            "seed": self.seed,
         }
+        return kwargs
 
     def create_module(
         self,
@@ -90,6 +99,18 @@ class MLPConfig(C.Config):
         )
 
 
+@contextlib.contextmanager
+def custom_seed_context(seed: int | None):
+    with contextlib.ExitStack() as stack:
+        if seed is not None:
+            stack.enter_context(
+                torch.random.fork_rng(devices=range(torch.cuda.device_count()))
+            )
+            torch.manual_seed(seed)
+
+        yield
+
+
 def MLP(
     dims: Sequence[int],
     activation: NonlinearityConfigBase
@@ -108,6 +129,7 @@ def MLP(
     pre_layers: Sequence[nn.Module] = [],
     post_layers: Sequence[nn.Module] = [],
     linear_cls: LinearModuleConstructor = nn.Linear,
+    seed: int | None = None,
 ):
     """
     Constructs a multi-layer perceptron (MLP) with the given dimensions and activation function.
@@ -123,52 +145,55 @@ def MLP(
         residual (bool, optional): Whether to use residual connections between layers. Defaults to False.
         pre_layers (Sequence[nn.Module], optional): List of layers to insert before the linear layers. Defaults to [].
         post_layers (Sequence[nn.Module], optional): List of layers to insert after the linear layers. Defaults to [].
+        linear_cls (LinearModuleConstructor, optional): Linear module constructor to use. Defaults to nn.Linear.
+        seed (int | None, optional): Random seed to use for initialization. If None, the default Torch behavior is used. Defaults to None.
 
     Returns:
         nn.Sequential: The constructed MLP.
     """
 
-    if activation is None:
-        activation = nonlinearity
-
-    if len(dims) < 2:
-        raise ValueError("mlp requires at least 2 dimensions")
-    if ln is True:
-        ln = "pre"
-    elif isinstance(ln, str) and ln not in ("pre", "post"):
-        raise ValueError("ln must be a boolean or 'pre' or 'post'")
-
-    layers: list[nn.Module] = []
-    if ln == "pre":
-        layers.append(nn.LayerNorm(dims[0]))
-
-    layers.extend(pre_layers)
-
-    for i in range(len(dims) - 1):
-        in_features = dims[i]
-        out_features = dims[i + 1]
-        bias_ = bias and not (no_bias_scalar and out_features == 1)
-        layers.append(linear_cls(in_features, out_features, bias=bias_))
-        if dropout is not None:
-            layers.append(nn.Dropout(dropout))
-        if i < len(dims) - 2:
-            match activation:
-                case NonlinearityConfigBase():
-                    layers.append(activation.create_module())
-                case nn.Module():
-                    # In this case, we create a deep copy of the module to avoid sharing parameters (if any).
-                    layers.append(copy.deepcopy(activation))
-                case Callable():
-                    layers.append(activation())
-                case _:
-                    raise ValueError(
-                        "Either `nonlinearity` or `activation` must be provided"
-                    )
-
-    layers.extend(post_layers)
-
-    if ln == "post":
-        layers.append(nn.LayerNorm(dims[-1]))
-
-    cls = ResidualSequential if residual else nn.Sequential
-    return cls(*layers)
+    with custom_seed_context(seed):
+        if activation is None:
+            activation = nonlinearity
+
+        if len(dims) < 2:
+            raise ValueError("mlp requires at least 2 dimensions")
+        if ln is True:
+            ln = "pre"
+        elif isinstance(ln, str) and ln not in ("pre", "post"):
+            raise ValueError("ln must be a boolean or 'pre' or 'post'")
+
+        layers: list[nn.Module] = []
+        if ln == "pre":
+            layers.append(nn.LayerNorm(dims[0]))
+
+        layers.extend(pre_layers)
+
+        for i in range(len(dims) - 1):
+            in_features = dims[i]
+            out_features = dims[i + 1]
+            bias_ = bias and not (no_bias_scalar and out_features == 1)
+            layers.append(linear_cls(in_features, out_features, bias=bias_))
+            if dropout is not None:
+                layers.append(nn.Dropout(dropout))
+            if i < len(dims) - 2:
+                match activation:
+                    case NonlinearityConfigBase():
+                        layers.append(activation.create_module())
+                    case nn.Module():
+                        # In this case, we create a deep copy of the module to avoid sharing parameters (if any).
+                        layers.append(copy.deepcopy(activation))
+                    case Callable():
+                        layers.append(activation())
+                    case _:
+                        raise ValueError(
+                            "Either `nonlinearity` or `activation` must be provided"
+                        )
+
+        layers.extend(post_layers)
+
+        if ln == "post":
+            layers.append(nn.LayerNorm(dims[-1]))
+
+        cls = ResidualSequential if residual else nn.Sequential
+        return cls(*layers)

nshtrainer-1.0.0b41/src/nshtrainer/nn/tests/test_mlp.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from typing import cast
+
+import pytest
+import torch
+
+from nshtrainer.nn.mlp import MLP
+
+
+def test_mlp_seed_reproducibility():
+    """Test that the seed parameter in MLP ensures reproducible weights."""
+
+    # Test dimensions
+    dims = [10, 20, 5]
+
+    # Create two MLPs with the same seed
+    seed1 = 42
+    mlp1 = MLP(dims, activation=torch.nn.ReLU(), seed=seed1)
+    mlp2 = MLP(dims, activation=torch.nn.ReLU(), seed=seed1)
+
+    # Create an MLP with a different seed
+    seed2 = 123
+    mlp3 = MLP(dims, activation=torch.nn.ReLU(), seed=seed2)
+
+    # Check first layer weights
+    layer1_weights1 = cast(torch.Tensor, mlp1[0].weight)
+    layer1_weights2 = cast(torch.Tensor, mlp2[0].weight)
+    layer1_weights3 = cast(torch.Tensor, mlp3[0].weight)
+
+    # Same seed should produce identical weights
+    assert torch.allclose(layer1_weights1, layer1_weights2)
+
+    # Different seeds should produce different weights
+    assert not torch.allclose(layer1_weights1, layer1_weights3)
+
+    # Check second layer weights
+    layer2_weights1 = cast(torch.Tensor, mlp1[2].weight)
+    layer2_weights2 = cast(torch.Tensor, mlp2[2].weight)
+    layer2_weights3 = cast(torch.Tensor, mlp3[2].weight)
+
+    # Same seed should produce identical weights for all layers
+    assert torch.allclose(layer2_weights1, layer2_weights2)
+
+    # Different seeds should produce different weights for all layers
+    assert not torch.allclose(layer2_weights1, layer2_weights3)
+
+    # Test that not providing a seed gives different results each time
+    mlp4 = MLP(dims, activation=torch.nn.ReLU(), seed=None)
+    mlp5 = MLP(dims, activation=torch.nn.ReLU(), seed=None)
+
+    # Without seeds, weights should be different
+    assert not torch.allclose(
+        cast(torch.Tensor, mlp4[0].weight), cast(torch.Tensor, mlp5[0].weight)
+    )