nshtrainer 0.6.0__tar.gz → 0.6.1__tar.gz

{nshtrainer-0.6.0 → nshtrainer-0.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.6.0
+Version: 0.6.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.6.0 → nshtrainer-0.6.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.6.0"
+version = "0.6.1"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

{nshtrainer-0.6.0 → nshtrainer-0.6.1}/src/nshtrainer/nn/__init__.py

@@ -1,4 +1,6 @@
 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 .module_dict import TypedModuleDict as TypedModuleDict
 from .module_list import TypedModuleList as TypedModuleList

{nshtrainer-0.6.0 → nshtrainer-0.6.1}/src/nshtrainer/nn/mlp.py

@@ -2,9 +2,10 @@ import copy
 from collections.abc import Callable, Sequence
 from typing import Literal, Protocol, runtime_checkable
 
+import nshconfig as C
 import torch
 import torch.nn as nn
-from typing_extensions import override
+from typing_extensions import TypedDict, override
 
 from .nonlinearity import BaseNonlinearityConfig
 
@@ -22,6 +23,71 @@ class ResidualSequential(nn.Sequential):
         return input + super().forward(input)
 
 
+class MLPConfigDict(TypedDict):
+    bias: bool
+    """Whether to include bias terms in the linear layers."""
+
+    no_bias_scalar: bool
+    """Whether to exclude bias terms when the output dimension is 1."""
+
+    nonlinearity: BaseNonlinearityConfig | None
+    """Activation function to use between layers."""
+
+    ln: bool | Literal["pre", "post"]
+    """Whether to apply layer normalization before or after the linear layers."""
+
+    dropout: float | None
+    """Dropout probability to apply between layers."""
+
+    residual: bool
+    """Whether to use residual connections between layers."""
+
+
+class MLPConfig(C.Config):
+    bias: bool = True
+    """Whether to include bias terms in the linear layers."""
+
+    no_bias_scalar: bool = True
+    """Whether to exclude bias terms when the output dimension is 1."""
+
+    nonlinearity: BaseNonlinearityConfig | None = None
+    """Activation function to use between layers."""
+
+    ln: bool | Literal["pre", "post"] = False
+    """Whether to apply layer normalization before or after the linear layers."""
+
+    dropout: float | None = None
+    """Dropout probability to apply between layers."""
+
+    residual: bool = False
+    """Whether to use residual connections between layers."""
+
+    def to_kwargs(self) -> MLPConfigDict:
+        return {
+            "bias": self.bias,
+            "no_bias_scalar": self.no_bias_scalar,
+            "nonlinearity": self.nonlinearity,
+            "ln": self.ln,
+            "dropout": self.dropout,
+            "residual": self.residual,
+        }
+
+    def create_module(
+        self,
+        dims: Sequence[int],
+        pre_layers: Sequence[nn.Module] = [],
+        post_layers: Sequence[nn.Module] = [],
+        linear_cls: LinearModuleConstructor = nn.Linear,
+    ):
+        return MLP(
+            dims,
+            **self.to_kwargs(),
+            pre_layers=pre_layers,
+            post_layers=post_layers,
+            linear_cls=linear_cls,
+        )
+
+
 def MLP(
     dims: Sequence[int],
     activation: BaseNonlinearityConfig