nshtrainer 0.6.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

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/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

{nshtrainer-0.6.0.dist-info → nshtrainer-0.6.1.dist-info}/METADATA

@@ -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.dist-info → nshtrainer-0.6.1.dist-info}/RECORD

@@ -41,8 +41,8 @@ nshtrainer/model/modules/logger.py,sha256=XEeo3QrplTNKZqfl6iWZf3fze3R4YOeOvs-RKV
 nshtrainer/model/modules/profiler.py,sha256=rQ_jRMcM1Z2AIROZlRnBRHM5rkTpq67afZPD6CIRfXs,825
 nshtrainer/model/modules/rlp_sanity_checks.py,sha256=o6gUceFwsuDHmL8eLOYuT3JGXFzq_qc4awl2RWaBygU,8900
 nshtrainer/model/modules/shared_parameters.py,sha256=mD5wrlBE3c025vzVdTpnSyC8yxzuI-aUWMmPhqPT0a0,2694
-nshtrainer/nn/__init__.py,sha256=57LPaP3G-BBGD2eGxbBUABNgYl3s_oASwrtOSS4bzTs,1339
-nshtrainer/nn/mlp.py,sha256=i-dHk0tomO_XlU6cKN4CC4HxTaYb-ukBCAgY1ySXl4I,3963
+nshtrainer/nn/__init__.py,sha256=0QPFl02a71WZQjLMGOlFNMmsYP5aa1q3eABHmnWH58Q,1427
+nshtrainer/nn/mlp.py,sha256=tX1VdtdzB0dyMkV0oEGCcob9hsYrXEIeusXfb-bJ5lQ,5940
 nshtrainer/nn/module_dict.py,sha256=NOY0B6WDTnktyWH4GthsprMQo0bpehC-hCq9SfD8paE,2329
 nshtrainer/nn/module_list.py,sha256=fb2u5Rqdjff8Pekyr9hkCPkBorQ-fldzzFAjsgWAm30,1719
 nshtrainer/nn/nonlinearity.py,sha256=owtU4kh4G98psD0axOJWVfBhm-OtJVgFM-TXSHmbNPU,3625
@@ -59,6 +59,6 @@ nshtrainer/util/seed.py,sha256=HEXgVs-wldByahOysKwq7506OHxdYTEgmP-tDQVAEkQ,287
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.6.0.dist-info/METADATA,sha256=o-z1r1fJTr1j-0Wv9USy61kD2VyALRjHST8sGAyLfBY,788
-nshtrainer-0.6.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.6.0.dist-info/RECORD,,
+nshtrainer-0.6.1.dist-info/METADATA,sha256=A99lygdq2iZY6oebqY5iOFcU71GX0t_cQN3cwKUlCVg,788
+nshtrainer-0.6.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.6.1.dist-info/RECORD,,