nshtrainer 1.4.1__py3-none-any.whl → 1.5.1__py3-none-any.whl

nshtrainer/_callback.py

@@ -8,38 +8,46 @@ from lightning.pytorch import LightningModule
 from lightning.pytorch.callbacks import Callback as _LightningCallback
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from torch.optim import Optimizer
+from typing_extensions import override
 
 if TYPE_CHECKING:
     from .trainer import Trainer
 
 
 class NTCallbackBase(_LightningCallback):
+    @override
     def setup(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule, stage: str
     ) -> None:
         """Called when fit, validate, test, predict, or tune begins."""
 
+    @override
     def teardown(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule, stage: str
     ) -> None:
         """Called when fit, validate, test, predict, or tune ends."""
 
+    @override
     def on_fit_start(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when fit begins."""
 
+    @override
     def on_fit_end(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when fit ends."""
 
+    @override
     def on_sanity_check_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the validation sanity check starts."""
 
+    @override
     def on_sanity_check_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the validation sanity check ends."""
 
+    @override
     def on_train_batch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
@@ -49,6 +57,7 @@ class NTCallbackBase(_LightningCallback):
     ) -> None:
         """Called when the train batch begins."""
 
+    @override
     def on_train_batch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
@@ -65,11 +74,13 @@ class NTCallbackBase(_LightningCallback):
 
         """
 
+    @override
     def on_train_epoch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the train epoch begins."""
 
+    @override
     def on_train_epoch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
@@ -81,10 +92,12 @@ class NTCallbackBase(_LightningCallback):
         .. code-block:: python
 
             class MyLightningModule(L.LightningModule):
+                 @override
                  def __init__(self):
                     super().__init__() # pyright: ignore[reportIncompatibleMethodOverride]
                     self.training_step_outputs = []
 
+                 @override
                  def training_step(self):
                     loss = ... # pyright: ignore[reportIncompatibleMethodOverride]
                     self.training_step_outputs.append(loss)
@@ -92,6 +105,7 @@ class NTCallbackBase(_LightningCallback):
 
 
             class MyCallback(L.Callback):
+                 @override
                  def on_train_epoch_end(self, trainer, pl_module):
                     # do something with all training_step outputs, for example: # pyright: ignore[reportIncompatibleMethodOverride]
                     epoch_mean = torch.stack(pl_module.training_step_outputs).mean()
@@ -101,36 +115,43 @@ class NTCallbackBase(_LightningCallback):
 
         """
 
+    @override
     def on_validation_epoch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the val epoch begins."""
 
+    @override
     def on_validation_epoch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the val epoch ends."""
 
+    @override
     def on_test_epoch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the test epoch begins."""
 
+    @override
     def on_test_epoch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the test epoch ends."""
 
+    @override
     def on_predict_epoch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the predict epoch begins."""
 
+    @override
     def on_predict_epoch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the predict epoch ends."""
 
+    @override
     def on_validation_batch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
@@ -141,6 +162,7 @@ class NTCallbackBase(_LightningCallback):
     ) -> None:
         """Called when the validation batch begins."""
 
+    @override
     def on_validation_batch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
@@ -152,6 +174,7 @@ class NTCallbackBase(_LightningCallback):
     ) -> None:
         """Called when the validation batch ends."""
 
+    @override
     def on_test_batch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
@@ -162,6 +185,7 @@ class NTCallbackBase(_LightningCallback):
     ) -> None:
         """Called when the test batch begins."""
 
+    @override
     def on_test_batch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
@@ -173,6 +197,7 @@ class NTCallbackBase(_LightningCallback):
     ) -> None:
         """Called when the test batch ends."""
 
+    @override
     def on_predict_batch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
@@ -183,6 +208,7 @@ class NTCallbackBase(_LightningCallback):
     ) -> None:
         """Called when the predict batch begins."""
 
+    @override
     def on_predict_batch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
@@ -194,36 +220,45 @@ class NTCallbackBase(_LightningCallback):
     ) -> None:
         """Called when the predict batch ends."""
 
+    @override
     def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when the train begins."""
 
+    @override
     def on_train_end(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when the train ends."""
 
+    @override
     def on_validation_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the validation loop begins."""
 
+    @override
     def on_validation_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the validation loop ends."""
 
+    @override
     def on_test_start(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when the test begins."""
 
+    @override
     def on_test_end(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when the test ends."""
 
+    @override
     def on_predict_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the predict begins."""
 
+    @override
     def on_predict_end(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when predict ends."""
 
+    @override
     def on_exception(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
@@ -232,7 +267,8 @@ class NTCallbackBase(_LightningCallback):
     ) -> None:
         """Called when any trainer execution is interrupted by an exception."""
 
-    def state_dict(self) -> dict[str, Any]:  # pyright: ignore[reportIncompatibleMethodOverride]
+    @override
+    def state_dict(self) -> dict[str, Any]:
         """Called when saving a checkpoint, implement to generate callback's ``state_dict``.
 
         Returns:
@@ -241,7 +277,8 @@ class NTCallbackBase(_LightningCallback):
         """
         return {}
 
-    def load_state_dict(self, state_dict: dict[str, Any]) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+    @override
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
         """Called when loading a checkpoint, implement to reload callback state given callback's ``state_dict``.
 
         Args:
@@ -250,6 +287,7 @@ class NTCallbackBase(_LightningCallback):
         """
         pass
 
+    @override
     def on_save_checkpoint(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
@@ -265,6 +303,7 @@ class NTCallbackBase(_LightningCallback):
 
         """
 
+    @override
     def on_load_checkpoint(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
@@ -280,16 +319,19 @@ class NTCallbackBase(_LightningCallback):
 
         """
 
+    @override
     def on_before_backward(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule, loss: torch.Tensor
     ) -> None:
         """Called before ``loss.backward()``."""
 
+    @override
     def on_after_backward(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called after ``loss.backward()`` and before optimizers are stepped."""
 
+    @override
     def on_before_optimizer_step(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
@@ -298,6 +340,7 @@ class NTCallbackBase(_LightningCallback):
     ) -> None:
         """Called before ``optimizer.step()``."""
 
+    @override
     def on_before_zero_grad(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
@@ -306,7 +349,10 @@ class NTCallbackBase(_LightningCallback):
     ) -> None:
         """Called before ``optimizer.zero_grad()``."""
 
-    def on_checkpoint_saved(  # pyright: ignore[reportIncompatibleMethodOverride]
+    # =================================================================
+    # Our own new callbacks
+    # =================================================================
+    def on_checkpoint_saved(
         self,
         ckpt_path: Path,
         metadata_path: Path | None,
@@ -317,6 +363,7 @@ class NTCallbackBase(_LightningCallback):
         pass
 
 
+@override
 def _call_on_checkpoint_saved(
     trainer: Trainer,
     ckpt_path: str | Path,

nshtrainer/callbacks/__init__.py

@@ -75,5 +75,5 @@ from .wandb_watch import WandbWatchCallbackConfig as WandbWatchCallbackConfig
 
 CallbackConfig = TypeAliasType(
     "CallbackConfig",
-    Annotated[CallbackConfigBase, callback_registry.DynamicResolution()],
+    Annotated[CallbackConfigBase, callback_registry],
 )

nshtrainer/callbacks/checkpoint/_base.py

@@ -5,13 +5,13 @@ import string
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, Literal
 
 import numpy as np
 import torch
 from lightning.pytorch import Trainer
 from lightning.pytorch.callbacks import Checkpoint
-from typing_extensions import override
+from typing_extensions import TypeVar, override
 
 from ..._checkpoint.metadata import CheckpointMetadata, _generate_checkpoint_metadata
 from ..._checkpoint.saver import link_checkpoint, remove_checkpoint

nshtrainer/callbacks/log_epoch.py

@@ -35,7 +35,7 @@ class LogEpochCallbackConfig(CallbackConfigBase):
         yield LogEpochCallback(self)
 
 
-def _worker_fn(
+def _log_on_step(
     trainer: Trainer,
     pl_module: LightningModule,
     num_batches_prop: str,
@@ -75,6 +75,19 @@ def _worker_fn(
     pl_module.log(metric_name, epoch, on_step=True, on_epoch=False)
 
 
+def _log_on_epoch(
+    trainer: Trainer,
+    pl_module: LightningModule,
+    *,
+    metric_name: str,
+):
+    if trainer.logger is None:
+        return
+
+    epoch = pl_module.current_epoch + 1
+    pl_module.log(metric_name, epoch, on_step=False, on_epoch=True)
+
+
 class LogEpochCallback(Callback):
     def __init__(self, config: LogEpochCallbackConfig):
         super().__init__()
@@ -85,16 +98,27 @@ class LogEpochCallback(Callback):
     def on_train_batch_start(
         self, trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int
     ):
-        if trainer.logger is None or not self.config.train:
+        if not self.config.train:
             return
 
-        _worker_fn(
+        _log_on_step(
             trainer,
             pl_module,
             "num_training_batches",
             metric_name=self.config.metric_name,
         )
 
+    @override
+    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.config.train:
+            return
+
+        _log_on_epoch(
+            trainer,
+            pl_module,
+            metric_name=self.config.metric_name,
+        )
+
     @override
     def on_validation_batch_start(
         self,
@@ -104,10 +128,10 @@ class LogEpochCallback(Callback):
         batch_idx: int,
         dataloader_idx: int = 0,
     ) -> None:
-        if trainer.logger is None or not self.config.val:
+        if not self.config.val:
             return
 
-        _worker_fn(
+        _log_on_step(
             trainer,
             pl_module,
             "num_val_batches",
@@ -115,6 +139,19 @@ class LogEpochCallback(Callback):
             metric_name=self.config.metric_name,
         )
 
+    @override
+    def on_validation_epoch_end(
+        self, trainer: Trainer, pl_module: LightningModule
+    ) -> None:
+        if not self.config.val:
+            return
+
+        _log_on_epoch(
+            trainer,
+            pl_module,
+            metric_name=self.config.metric_name,
+        )
+
     @override
     def on_test_batch_start(
         self,
@@ -124,13 +161,24 @@ class LogEpochCallback(Callback):
         batch_idx: int,
         dataloader_idx: int = 0,
     ) -> None:
-        if trainer.logger is None or not self.config.test:
+        if not self.config.test:
             return
 
-        _worker_fn(
+        _log_on_step(
             trainer,
             pl_module,
             "num_test_batches",
             dataloader_idx=dataloader_idx,
             metric_name=self.config.metric_name,
         )
+
+    @override
+    def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        if not self.config.test:
+            return
+
+        _log_on_epoch(
+            trainer,
+            pl_module,
+            metric_name=self.config.metric_name,
+        )

nshtrainer/callbacks/print_table.py

@@ -67,14 +67,14 @@ class PrintTableMetricsCallback(Callback):
             }
         self.metrics.append(metrics_dict)
 
-        from rich.console import Console  # type: ignore[reportMissingImports] # noqa
+        from rich.console import Console  # pyright: ignore[reportMissingImports] # noqa
 
         console = Console()
         table = self.create_metrics_table()
         console.print(table)
 
     def create_metrics_table(self):
-        from rich.table import Table  # type: ignore[reportMissingImports] # noqa
+        from rich.table import Table  # pyright: ignore[reportMissingImports] # noqa
 
         table = Table(show_header=True, header_style="bold magenta")
 

nshtrainer/callbacks/rlp_sanity_checks.py

@@ -38,6 +38,7 @@ class RLPSanityChecksCallbackConfig(CallbackConfigBase):
     def __bool__(self):
         return self.enabled
 
+    @override
     def create_callbacks(self, trainer_config):
         if not self:
             return

nshtrainer/configs/__init__.py

@@ -111,7 +111,6 @@ from nshtrainer.optimizer import RAdamConfig as RAdamConfig
 from nshtrainer.optimizer import RMSpropConfig as RMSpropConfig
 from nshtrainer.optimizer import RpropConfig as RpropConfig
 from nshtrainer.optimizer import SGDConfig as SGDConfig
-from nshtrainer.optimizer import Union as Union
 from nshtrainer.optimizer import optimizer_registry as optimizer_registry
 from nshtrainer.profiler import AdvancedProfilerConfig as AdvancedProfilerConfig
 from nshtrainer.profiler import BaseProfilerConfig as BaseProfilerConfig
@@ -355,7 +354,6 @@ __all__ = [
     "TorchSyncBatchNormPlugin",
     "TrainerConfig",
     "TransformerEnginePluginConfig",
-    "Union",
     "WandbLoggerConfig",
     "WandbUploadCodeCallbackConfig",
     "WandbWatchCallbackConfig",

nshtrainer/configs/optimizer/__init__.py

@@ -16,7 +16,6 @@ from nshtrainer.optimizer import RAdamConfig as RAdamConfig
 from nshtrainer.optimizer import RMSpropConfig as RMSpropConfig
 from nshtrainer.optimizer import RpropConfig as RpropConfig
 from nshtrainer.optimizer import SGDConfig as SGDConfig
-from nshtrainer.optimizer import Union as Union
 from nshtrainer.optimizer import optimizer_registry as optimizer_registry
 
 __all__ = [
@@ -34,6 +33,5 @@ __all__ = [
     "RMSpropConfig",
     "RpropConfig",
     "SGDConfig",
-    "Union",
     "optimizer_registry",
 ]

nshtrainer/loggers/__init__.py

@@ -12,6 +12,5 @@ from .tensorboard import TensorboardLoggerConfig as TensorboardLoggerConfig
 from .wandb import WandbLoggerConfig as WandbLoggerConfig
 
 LoggerConfig = TypeAliasType(
-    "LoggerConfig",
-    Annotated[LoggerConfigBase, logger_registry.DynamicResolution()],
+    "LoggerConfig", Annotated[LoggerConfigBase, logger_registry]
 )

nshtrainer/loggers/actsave.py

@@ -5,7 +5,7 @@ from typing import Any, Literal
 
 import numpy as np
 from lightning.pytorch.loggers import Logger
-from typing_extensions import final
+from typing_extensions import final, override
 
 from .base import LoggerConfigBase, logger_registry
 
@@ -15,6 +15,7 @@ from .base import LoggerConfigBase, logger_registry
 class ActSaveLoggerConfig(LoggerConfigBase):
     name: Literal["actsave"] = "actsave"
 
+    @override
     def create_logger(self, trainer_config):
         if not self.enabled:
             return None
@@ -24,10 +25,12 @@ class ActSaveLoggerConfig(LoggerConfigBase):
 
 class ActSaveLogger(Logger):
     @property
+    @override
     def name(self):
         return None
 
     @property
+    @override
     def version(self):
         from nshutils import ActSave
 
@@ -37,6 +40,7 @@ class ActSaveLogger(Logger):
         return ActSave._saver._id
 
     @property
+    @override
     def save_dir(self):
         from nshutils import ActSave
 
@@ -45,6 +49,7 @@ class ActSaveLogger(Logger):
 
         return str(ActSave._saver._save_dir)
 
+    @override
     def log_hyperparams(
         self,
         params: dict[str, Any] | Namespace,
@@ -56,6 +61,7 @@ class ActSaveLogger(Logger):
         # Wrap the hparams as a object-dtype np array
         return ActSave.save({"hyperparameters": np.array(params, dtype=object)})
 
+    @override
     def log_metrics(self, metrics: dict[str, float], step: int | None = None) -> None:
         from nshutils import ActSave
 

nshtrainer/loggers/wandb.py

@@ -63,7 +63,7 @@ class FinishWandbOnTeardownCallback(Callback):
         stage: str,
     ):
         try:
-            import wandb  # type: ignore
+            import wandb
         except ImportError:
             return
 
@@ -139,7 +139,7 @@ class WandbLoggerConfig(CallbackConfigBase, LoggerConfigBase):
         # If `wandb-core` is enabled, we should use the new backend.
         if self.use_wandb_core:
             try:
-                import wandb  # type: ignore
+                import wandb
 
                 # The minimum version that supports the new backend is 0.17.5
                 wandb_version = version.parse(importlib.metadata.version("wandb"))
@@ -151,7 +151,7 @@ class WandbLoggerConfig(CallbackConfigBase, LoggerConfigBase):
                     )
                 # W&B versions 0.18.0 use wandb-core by default
                 elif wandb_version < version.parse("0.18.0"):
-                    wandb.require("core")  # type: ignore
+                    wandb.require("core")
                     log.critical("Using the `wandb-core` backend for WandB.")
             except ImportError:
                 pass
@@ -166,9 +166,9 @@ class WandbLoggerConfig(CallbackConfigBase, LoggerConfigBase):
                     "If you want to use the new `wandb-core` backend, set `use_wandb_core=True`."
                 )
                 try:
-                    import wandb  # type: ignore
+                    import wandb
 
-                    wandb.require("legacy-service")  # type: ignore
+                    wandb.require("legacy-service")
                 except ImportError:
                     pass
 

nshtrainer/lr_scheduler/base.py

@@ -81,7 +81,7 @@ class LRSchedulerConfigBase(C.Config, ABC):
             scheduler["monitor"] = metadata["monitor"]
         # - `strict`
         if scheduler.get("strict") is None and "strict" in metadata:
-            scheduler["strict"] = metadata["strict"]  # type: ignore
+            scheduler["strict"] = metadata["strict"]
 
         return scheduler
 

nshtrainer/model/mixins/callback.py

@@ -41,23 +41,6 @@ class CallbackModuleMixin(
     CallbackRegistrarModuleMixin,
     mixin_base_type(LightningModule),
 ):
-    @property
-    def _nshtrainer_callbacks(self) -> list[CallbackFn]:
-        if not hasattr(self, "_private_nshtrainer_callbacks_list"):
-            self._private_nshtrainer_callbacks_list = []
-        return self._private_nshtrainer_callbacks_list
-
-    def register_callback(
-        self,
-        callback: _Callback | Iterable[_Callback] | CallbackFn | None = None,
-    ):
-        if not callable(callback):
-            callback_ = cast(CallbackFn, lambda: callback)
-        else:
-            callback_ = callback
-
-        self._nshtrainer_callbacks.append(callback_)
-
     def _gather_all_callbacks(self):
         modules: list[Any] = []
         if isinstance(self, CallbackRegistrarModuleMixin):

nshtrainer/model/mixins/logger.py

@@ -203,6 +203,7 @@ class LoggerLightningModuleMixin(mixin_base_type(LightningModule)):
         name = f"{prefix}{name}"
         return super().log(name, value, metric_attribute=metric_attribute, **fn_kwargs)
 
+    @override
     def log_dict(
         self,
         dictionary: Mapping[str, _METRIC] | torchmetrics.MetricCollection,

nshtrainer/nn/module_dict.py

@@ -28,9 +28,9 @@ class TypedModuleDict(nn.Module, Generic[TModule]):
         return f"{self.key_prefix}{key}"
 
     def _remove_prefix(self, key: str) -> str:
-        assert key.startswith(
-            self.key_prefix
-        ), f"{key} does not start with {self.key_prefix}"
+        assert key.startswith(self.key_prefix), (
+            f"{key} does not start with {self.key_prefix}"
+        )
         return key[len(self.key_prefix) :]
 
     def __setitem__(self, key: str, module: TModule) -> None:
@@ -39,7 +39,7 @@ class TypedModuleDict(nn.Module, Generic[TModule]):
 
     def __getitem__(self, key: str) -> TModule:
         key = self._with_prefix(key)
-        return self._module_dict.__getitem__(key)  # type: ignore
+        return cast(TModule, self._module_dict.__getitem__(key))
 
     def update(self, modules: Mapping[str, TModule]) -> None:
         return self._module_dict.update(

nshtrainer/nn/module_list.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
 from collections.abc import Iterable, Iterator
-from typing import Generic, TypeVar, overload
+from typing import Generic, cast, overload
 
 import torch.nn as nn
-from typing_extensions import override
+from typing_extensions import TypeVar, override
 
-TModule = TypeVar("TModule", bound=nn.Module)
+TModule = TypeVar("TModule", bound=nn.Module, infer_variance=True)
 
 
 class TypedModuleList(nn.ModuleList, Generic[TModule]):
@@ -14,39 +14,39 @@ class TypedModuleList(nn.ModuleList, Generic[TModule]):
         super().__init__(modules)
 
     @overload
-    def __getitem__(self, idx: slice) -> "TypedModuleList[TModule]": ...
+    def __getitem__(self, idx: slice) -> TypedModuleList[TModule]: ...
 
     @overload
     def __getitem__(self, idx: int) -> TModule: ...
 
     @override
-    def __getitem__(self, idx: int | slice) -> TModule | "TypedModuleList[TModule]":
-        return super().__getitem__(idx)  # type: ignore
+    def __getitem__(self, idx: int | slice) -> TModule | TypedModuleList[TModule]:
+        return cast(TModule | TypedModuleList[TModule], super().__getitem__(idx))
 
     @override
-    def __setitem__(self, idx: int, module: TModule) -> None:  # type: ignore
+    def __setitem__(self, idx: int, module: TModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         return super().__setitem__(idx, module)
 
     @override
     def __iter__(self) -> Iterator[TModule]:
-        return super().__iter__()  # type: ignore
+        return cast(Iterator[TModule], super().__iter__())
 
     @override
-    def __iadd__(self, modules: Iterable[TModule]) -> "TypedModuleList[TModule]":  # type: ignore
-        return super().__iadd__(modules)  # type: ignore
+    def __iadd__(self, modules: Iterable[TModule]) -> TypedModuleList[TModule]:  # pyright: ignore[reportIncompatibleMethodOverride]
+        return cast(TypedModuleList[TModule], super().__iadd__(modules))
 
     @override
-    def __add__(self, modules: Iterable[TModule]) -> "TypedModuleList[TModule]":  # type: ignore
-        return super().__add__(modules)  # type: ignore
+    def __add__(self, modules: Iterable[TModule]) -> TypedModuleList[TModule]:  # pyright: ignore[reportIncompatibleMethodOverride]
+        return cast(TypedModuleList[TModule], super().__add__(modules))
 
     @override
-    def insert(self, idx: int, module: TModule) -> None:  # type: ignore
+    def insert(self, idx: int, module: TModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         return super().insert(idx, module)
 
     @override
-    def append(self, module: TModule) -> "TypedModuleList[TModule]":  # type: ignore
-        return super().append(module)  # type: ignore
+    def append(self, module: TModule) -> TypedModuleList[TModule]:  # pyright: ignore[reportIncompatibleMethodOverride]
+        return cast(TypedModuleList[TModule], super().append(module))
 
     @override
-    def extend(self, modules: Iterable[TModule]) -> "TypedModuleList[TModule]":  # type: ignore
-        return super().extend(modules)  # type: ignore
+    def extend(self, modules: Iterable[TModule]) -> TypedModuleList[TModule]:  # pyright: ignore[reportIncompatibleMethodOverride]
+        return cast(TypedModuleList[TModule], super().extend(modules))

nshtrainer/nn/nonlinearity.py

@@ -30,6 +30,7 @@ class ReLUNonlinearityConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return nn.ReLU()
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         return F.relu(x)
 
@@ -43,6 +44,7 @@ class SigmoidNonlinearityConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return nn.Sigmoid()
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         return torch.sigmoid(x)
 
@@ -56,6 +58,7 @@ class TanhNonlinearityConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return nn.Tanh()
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         return torch.tanh(x)
 
@@ -72,6 +75,7 @@ class SoftmaxNonlinearityConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return nn.Softmax(dim=self.dim)
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         return torch.softmax(x, dim=self.dim)
 
@@ -91,6 +95,7 @@ class SoftplusNonlinearityConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return nn.Softplus(beta=self.beta, threshold=self.threshold)
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         return F.softplus(x, beta=self.beta, threshold=self.threshold)
 
@@ -104,6 +109,7 @@ class SoftsignNonlinearityConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return nn.Softsign()
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         return F.softsign(x)
 
@@ -120,6 +126,7 @@ class ELUNonlinearityConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return nn.ELU()
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         return F.elu(x, alpha=self.alpha)
 
@@ -136,6 +143,7 @@ class LeakyReLUNonlinearityConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return nn.LeakyReLU(negative_slope=self.negative_slope)
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         return F.leaky_relu(x, negative_slope=self.negative_slope)
 
@@ -157,6 +165,7 @@ class PReLUConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return nn.PReLU(num_parameters=self.num_parameters, init=self.init)
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         raise NotImplementedError(
             "PReLU requires learnable parameters and cannot be called directly."
@@ -175,6 +184,7 @@ class GELUNonlinearityConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return nn.GELU(approximate=self.approximate)
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         return F.gelu(x, approximate=self.approximate)
 
@@ -188,6 +198,7 @@ class SwishNonlinearityConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return nn.SiLU()
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         return F.silu(x)
 
@@ -201,6 +212,7 @@ class SiLUNonlinearityConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return nn.SiLU()
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         return F.silu(x)
 
@@ -214,6 +226,7 @@ class MishNonlinearityConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return nn.Mish()
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         return F.mish(x)
 
@@ -234,12 +247,12 @@ class SwiGLUNonlinearityConfig(NonlinearityConfigBase):
     def create_module(self) -> nn.Module:
         return SwiGLU()
 
+    @override
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         input, gate = x.chunk(2, dim=-1)
         return input * F.silu(gate)
 
 
 NonlinearityConfig = TypeAliasType(
-    "NonlinearityConfig",
-    Annotated[NonlinearityConfigBase, nonlinearity_registry.DynamicResolution()],
+    "NonlinearityConfig", Annotated[NonlinearityConfigBase, nonlinearity_registry]
 )

nshtrainer/optimizer.py

@@ -2,11 +2,10 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
-from typing import Annotated, Any, Literal, Tuple, Union
+from typing import Annotated, Any, Literal
 
 import nshconfig as C
 import torch.nn as nn
-from torch import Tensor
 from torch.optim import Optimizer
 from typing_extensions import TypeAliasType, final, override
 
@@ -621,6 +620,5 @@ class SGDConfig(OptimizerConfigBase):
 
 
 OptimizerConfig = TypeAliasType(
-    "OptimizerConfig",
-    Annotated[OptimizerConfigBase, optimizer_registry.DynamicResolution()],
+    "OptimizerConfig", Annotated[OptimizerConfigBase, optimizer_registry]
 )

nshtrainer/trainer/accelerator.py

@@ -23,8 +23,7 @@ class AcceleratorConfigBase(C.Config, ABC):
 accelerator_registry = C.Registry(AcceleratorConfigBase, discriminator="name")
 
 AcceleratorConfig = TypeAliasType(
-    "AcceleratorConfig",
-    Annotated[AcceleratorConfigBase, accelerator_registry.DynamicResolution()],
+    "AcceleratorConfig", Annotated[AcceleratorConfigBase, accelerator_registry]
 )
 
 

nshtrainer/trainer/plugin/__init__.py

@@ -13,6 +13,5 @@ from .base import PluginConfigBase as PluginConfigBase
 from .base import plugin_registry as plugin_registry
 
 PluginConfig = TypeAliasType(
-    "PluginConfig",
-    Annotated[PluginConfigBase, plugin_registry.DynamicResolution()],
+    "PluginConfig", Annotated[PluginConfigBase, plugin_registry]
 )

nshtrainer/util/code_upload.py

@@ -17,7 +17,7 @@ def get_code_dir() -> Path | None:
         # New versions of nshrunner will have the code_dir attribute
         #   in the session object. We should use that. Otherwise, use snapshot_dir.
         try:
-            code_dir = session.code_dir  # type: ignore
+            code_dir = session.code_dir
         except AttributeError:
             code_dir = session.snapshot_dir
 

{nshtrainer-1.4.1.dist-info → nshtrainer-1.5.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.4.1
+Version: 1.5.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.4.1.dist-info → nshtrainer-1.5.1.dist-info}/RECORD

@@ -1,15 +1,15 @@
 nshtrainer/.nshconfig.generated.json,sha256=yZd6cn1RhvNNJUgiUTRYut8ofZYvbulnpPG-rZIRhi4,106
 nshtrainer/__init__.py,sha256=RI_2B_IUWa10B6H5TAuWtE5FWX1X4ue-J4dTDaF2-lQ,1035
-nshtrainer/_callback.py,sha256=ZDppiJ4d65tRXTEWYPZLH_F1xFizdz1pkWJe_sQ5uII,12564
+nshtrainer/_callback.py,sha256=aBg9Za6hjteHcGjb8bIGzaN57A03cXrPv4rMWqaNsLU,13253
 nshtrainer/_checkpoint/metadata.py,sha256=El9Ip8jGA7mAN5rAMpVfg1dfUe2dGoOOfvF1JfYJGHM,5676
 nshtrainer/_checkpoint/saver.py,sha256=utcrYKSosd04N9m2GIylufO5DO05D90qVU3mvadfApU,1658
 nshtrainer/_experimental/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 nshtrainer/_hf_hub.py,sha256=OB4252GJ6AbKNCRmHVvEglvjYVMUN822BFYECABxfZU,14037
-nshtrainer/callbacks/__init__.py,sha256=m6eJuprZfBELuKpngKXre33B9yPXkG7jlKVmI-0yXRQ,4000
+nshtrainer/callbacks/__init__.py,sha256=6l2vrFywWftzKTlZMEkF-WgE5uLjLgX89BoUMq8_x-0,3980
 nshtrainer/callbacks/actsave.py,sha256=NSXIIu62MNYe5gz479SMW33bdoKYoYtWtd_iTWFpKpc,3881
 nshtrainer/callbacks/base.py,sha256=K9aom1WVVRYxl-tHWgtmDUQZ1o63NgznvLsjauTKcCc,4225
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=l8tkHc83_mLiU0-wT09SWdRzwpm2ulbkLzcuCmuTwzE,620
-nshtrainer/callbacks/checkpoint/_base.py,sha256=BjgfCXsf4Ihf1MNKkHBUwjHMLwc04PZO-2Bx-LdAazg,11010
+nshtrainer/callbacks/checkpoint/_base.py,sha256=AsNt1bZ-yloPHqenRy4KAJK5DDmhBY1RprR2_xbvomc,11010
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=aCs3E1eucfDlUeW2Iq_Ke7hb96BxHanmvn7PCCbqq0E,2648
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=vn-as3ex7kaTRcKsIurVtM6kUSHYNwHJeYG82j2dMcc,3554
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=nljzETqkHwA-4g8mxaeFK5HxA8My0dlIPzIUscSMWyk,3525
@@ -21,18 +21,18 @@ nshtrainer/callbacks/ema.py,sha256=dBFiUXG0xmyCw8-ayuSzJMKqSbepl6Ii5VIbhFlT5ug,1
 nshtrainer/callbacks/finite_checks.py,sha256=3lZ3kEIjmYQfqTF0DcrgZ9_98ZLQhQj8usH7SgWst3o,2185
 nshtrainer/callbacks/gradient_skipping.py,sha256=8g7oC7PF0LTAEzwiNoaS5tWOnkjk_EB0QG3JdHkQ8ek,3523
 nshtrainer/callbacks/interval.py,sha256=UCzUzt3XCFVyQyCWL9lOrStkkxesvduNOYk8yMrGTTk,8116
-nshtrainer/callbacks/log_epoch.py,sha256=C2yUww8lAuCX-dy06tsw95yCBOfFd2mfGs0VhrEq1oU,3775
+nshtrainer/callbacks/log_epoch.py,sha256=-uC5ss9p_ngXUCrSIUwViFcaaVX6ALUzIAKxoDgZrac,4823
 nshtrainer/callbacks/lr_monitor.py,sha256=v45ehnwNO987087HfiOY5aIrVRbwdKMgPYRFHs1fyEE,1444
 nshtrainer/callbacks/metric_validation.py,sha256=4RDr1FuNKfro-6QEtmcFqT4iNf2twmJVNk9y-8nq9bg,2882
 nshtrainer/callbacks/norm_logging.py,sha256=nVIDWe-ASl5zN830-ODR8QMCqI1ma-QPCIwoy0Wb-Nk,6390
-nshtrainer/callbacks/print_table.py,sha256=VaS4JgI963do79laXK4lUkFQx8v6aRSy22W0zyal_LA,3035
-nshtrainer/callbacks/rlp_sanity_checks.py,sha256=Df9Prq2QKXnaeMBIvMQBhDhJTDeru5UbiuXJOJR16Gk,10050
+nshtrainer/callbacks/print_table.py,sha256=xdDvogpLFHdaHM4yDGENvJUX4Gz4hDq-QpsPcv-Oqi8,3041
+nshtrainer/callbacks/rlp_sanity_checks.py,sha256=PRtcj9K9fa2Oh6nbKQJR6w2__on0Jln969qZXlnkv1Q,10064
 nshtrainer/callbacks/shared_parameters.py,sha256=s94jJTAIbDGukYJu6l247QonVOCudGClU4t5kLt8XrY,3076
 nshtrainer/callbacks/timer.py,sha256=gDcw_K_ikf0bkVgxQ0cDhvvNvz6GLZVLcatuKfh0ORU,4731
 nshtrainer/callbacks/wandb_upload_code.py,sha256=4X-mpiX5ghj9vnEreK2i8Xyvimqt0K-PNWA2HtT-B6I,1940
 nshtrainer/callbacks/wandb_watch.py,sha256=VB14Dy5ZRXQ3di0fPv0K_DFJurLhroLPytnuwQBiJFg,3037
 nshtrainer/configs/.gitattributes,sha256=VeZmarvNEqiRBOHGcllpKm90nL6C8u4tBu7SEm7fj-E,26
-nshtrainer/configs/__init__.py,sha256=-yJ5Uk9VkANqfk-QnX2aynL0jSf7cJQuQNzT1GAE1x8,15684
+nshtrainer/configs/__init__.py,sha256=ZHV_1zCZKUYBKzWiLPrF8eFKsb-gepAF4G7AmsCxkkA,15623
 nshtrainer/configs/_checkpoint/__init__.py,sha256=6s7Y68StboqscY2G4P_QG443jz5aiym5SjOogIljWLg,342
 nshtrainer/configs/_checkpoint/metadata/__init__.py,sha256=oOPfYkXTjKgm6pluGsG6V1TPyCEGjsQpHVL-LffSUFQ,290
 nshtrainer/configs/_hf_hub/__init__.py,sha256=ciFLbV-JV8SVzqo2SyythEuDMnk7gGfdIacB18QYnkY,511
@@ -77,7 +77,7 @@ nshtrainer/configs/nn/__init__.py,sha256=Ms2gIqbRxNVm6GHKCddCJTTqMwUPifjjHD_fCfJ
 nshtrainer/configs/nn/mlp/__init__.py,sha256=O6kQ6utZNJPG9Fax5pRdZcHa3J-XFKKdXcc_PQg0jk0,347
 nshtrainer/configs/nn/nonlinearity/__init__.py,sha256=LCTbTyelCMABVw505CGQ4UpEGlAnIhflSLFqwAQXLQA,2155
 nshtrainer/configs/nn/rng/__init__.py,sha256=4iC6vwxbfNeXyvpwZ1Z5Kcy-he4cu7mg3UpLD-RLrHc,141
-nshtrainer/configs/optimizer/__init__.py,sha256=8ztp5UD-edfzwF-qdJTeZwlv-YWJ5Sn230b9aWxJyQQ,1398
+nshtrainer/configs/optimizer/__init__.py,sha256=Kq6ACztSQhwgE_tP4F1RI7nQMBgC1ebQmY3HaBYKbeg,1337
 nshtrainer/configs/profiler/__init__.py,sha256=2ssaIpfVnvcbfNvZ-JeKp1Cx4NO1LknkVqTm1hu7Lvw,768
 nshtrainer/configs/profiler/_base/__init__.py,sha256=ekYfPg-VDhCAFM5nJka2TxUYdRDm1CKqjwUOQNbQjD4,176
 nshtrainer/configs/profiler/advanced/__init__.py,sha256=-ThpUat16Ij_0avkMUVVA8wCWDG_q_tM7KQofnWQCtg,308
@@ -103,30 +103,30 @@ nshtrainer/data/__init__.py,sha256=K4i3Tw4g9EOK2zlMMbidi99y0SyI4w8P7_XUf1n42Ts,2
 nshtrainer/data/balanced_batch_sampler.py,sha256=r1cBKRXKHD8E1Ax6tj-FUbE-z1qpbO58mQ9VrK9uLnc,5481
 nshtrainer/data/datamodule.py,sha256=Rb4-mA8iXtjRlNUHcIqVPEvxA_VkiJXwN1EvHIsydJ0,4095
 nshtrainer/data/transform.py,sha256=qd0lIocO59Fk_m90xyOHgFezbymd1mRwly8nbYIfHGc,2263
-nshtrainer/loggers/__init__.py,sha256=fI0OHEltHP4tZI-KFB3npdzoxm_M2QsEYKxY3um05_s,592
-nshtrainer/loggers/actsave.py,sha256=wgNrpBB6wQM7qff8iLDb_sQnbiAcYHRmH56pcEJPB3o,1409
+nshtrainer/loggers/__init__.py,sha256=0fnclaEIgAUrRlYuSKfzni11dlJ6edllrs06NVmbtYc,567
+nshtrainer/loggers/actsave.py,sha256=Xd21jaBVUmkxITKYfycWKZEgcHu1-dmi1H5EYEjvPDw,1503
 nshtrainer/loggers/base.py,sha256=ON92XbwTSgadQOSyw5PiRRFzyH6uJ-xLtE0nB3cbgPc,1205
 nshtrainer/loggers/csv.py,sha256=xJ8mSmw4vJwinIfqhF6t2HWmh_1dXEYyLfGuXwL7WHo,1160
 nshtrainer/loggers/tensorboard.py,sha256=E7iO_fDt9bfH02hBL430bXPLljOo5iGgq2QyPqmx2gQ,2324
-nshtrainer/loggers/wandb.py,sha256=KZXAUWrrmdX_L8rqej77oUHaM0JxZRM8y9z6JP9PISw,6856
+nshtrainer/loggers/wandb.py,sha256=EK2rvJwmV-LxXIm21ZORNBI9nz-AXqo2mIN7xyjs8bc,6776
 nshtrainer/lr_scheduler/__init__.py,sha256=daMMK3erUcNXGGd_nZB8AWu3ZTYqfS1RSWeK4FV2udw,851
-nshtrainer/lr_scheduler/base.py,sha256=LE53JRBTuAlA1fqbMgCZ7m39D1z0rGj2TizhJ62CPvE,3756
+nshtrainer/lr_scheduler/base.py,sha256=24I3PNlj2iYPoaHeD2_InMAptclCtMnZoD8nXWqxLYw,3740
 nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=MsoXgCcWTKsrkNZiGnKS6yC-slRuleuwFxeM_lmG_pQ,5560
 nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=irPyDjfUX843ze4bJM9sW8WSeEcU643QJ30JN2hz9Rc,3206
 nshtrainer/metrics/__init__.py,sha256=Nqkn_jsDf3n5WtfMcnaaEftYjIIT2b-S7rmsB1MOMkU,86
 nshtrainer/metrics/_config.py,sha256=ox_ScK6V0J9nzIMhEB0qpToNKpt83VVgOVSRFCV-wBc,595
 nshtrainer/model/__init__.py,sha256=3G-bwPPSRStWdsdwG9-rn0bXcRpEiP1BiQpF_qavtls,97
 nshtrainer/model/base.py,sha256=PvTmupfGahEZME0BWqbeErDPP1VOm2Nm9JxJkO8afcc,10815
-nshtrainer/model/mixins/callback.py,sha256=0LPgve4VszHbLipid4mpI1qnnmdGS2spivs0dXLvqHw,3154
+nshtrainer/model/mixins/callback.py,sha256=-walDV3fxH4K-ezugvL__Tml9OP1WIlaaTT8j6mWxLI,2580
 nshtrainer/model/mixins/debug.py,sha256=ydLuAAaa7M5bX0gougZ5gWuZnvn4Ra9assal3IZ9hq8,2086
-nshtrainer/model/mixins/logger.py,sha256=7u9fQig-SVFA9RFIB4U0gqJAzruh49mgmXXvZ6VkDUk,11694
+nshtrainer/model/mixins/logger.py,sha256=-D4YwSg0eTDtXj3N288FEo6rqsZ518u1aMBE4Dv4tmg,11708
 nshtrainer/nn/__init__.py,sha256=Vd246v2N9tBQ8XxmTquWzj5lAmeSnngrjpYOfp4LTXM,1499
 nshtrainer/nn/mlp.py,sha256=nYUgAISzuhC8sav6PloAdyz0PdEoikwppiXIuToEVdE,7550
-nshtrainer/nn/module_dict.py,sha256=9plb8aQUx5TUEPhX5jI9u8LrpTeKe7jZAHi8iIqcN8w,2365
-nshtrainer/nn/module_list.py,sha256=UB43pcwD_3nUke_DyLQt-iXKhWdKM6Zjm84lRC1hPYA,1755
-nshtrainer/nn/nonlinearity.py,sha256=xmaL4QCRvCxqmaGIOwetJeKK-6IK4m2OV7D3SjxSwJQ,6322
+nshtrainer/nn/module_dict.py,sha256=FJrxUgQkY6O6tmA_7I_kRoPvxLtPU3cYZY-42InVG3A,2366
+nshtrainer/nn/module_list.py,sha256=xvoF8F-pG-z3MnYc91anG9vUQFVes6niy8J8J0qVAlg,2091
+nshtrainer/nn/nonlinearity.py,sha256=UhAsc8o_6AIsos6sktUWC9xLFCFgHJn5WiurKN1sf5U,6493
 nshtrainer/nn/rng.py,sha256=IJGvX9v8qBkfgBrMlNU2aj-MbYTPoncFyJzvPkzCQpM,512
-nshtrainer/optimizer.py,sha256=8pjOny7NxIt04PXxn3zOyJ2soL7nmj8yBVV82r_tNsc,17522
+nshtrainer/optimizer.py,sha256=hvw_UNovYgLHhDvMr9BbUz3EPOIrGZDz9ir8lvCgiw0,17458
 nshtrainer/profiler/__init__.py,sha256=RjaNBoVcTFu8lF0dNlFp-2LaPYdonoIbDy2_KhgF0Ek,594
 nshtrainer/profiler/_base.py,sha256=kFcSVn9gJuMwgDxbfyHh46CmEAIPZjxw3yjPbKgzvwA,950
 nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5NOQ,1160
@@ -137,8 +137,8 @@ nshtrainer/trainer/_config.py,sha256=GL8DtuH-6x2aHcRlEcmzyhEBMRRldiSazNAeNmPw7gM
 nshtrainer/trainer/_distributed_prediction_result.py,sha256=bQw8Z6PT694UUf-zQPkech6CxyUSy8bAIexfSfPej0U,2507
 nshtrainer/trainer/_log_hparams.py,sha256=XH2lZ4U_3AZBhOt91ocsEhdL_NRz35oWvqLCUFDohUs,2389
 nshtrainer/trainer/_runtime_callback.py,sha256=6F2Gq27Q8OFfN3RtdNC6QRA8ac0LC1hh4DUE3V5WgbI,4217
-nshtrainer/trainer/accelerator.py,sha256=Bqq-ry7DeCY4zw9_zBvTZiijpA-uUHrDjtbLV652m4M,2415
-nshtrainer/trainer/plugin/__init__.py,sha256=q_q98MYNaZ2VE_tqGqYlQjQnlaF4NE1FUqVVbj0EK7k,517
+nshtrainer/trainer/accelerator.py,sha256=rWfSJ-pQsLREaRPF_rRXsqxvaQQf6XGT6zpNt829jk0,2390
+nshtrainer/trainer/plugin/__init__.py,sha256=LSxEK0vnoN9WkU8MDIetrVrDPLCowGLoc9cvh6RG6gg,492
 nshtrainer/trainer/plugin/base.py,sha256=76ct2TTHLpPr5MO8B9CIkoCOo-dFImzqAll8cIdC0cg,736
 nshtrainer/trainer/plugin/environment.py,sha256=SSXRWHjyFUA6oFx3duD_ZwhM59pWUjR1_UzHz02NI2c,5440
 nshtrainer/trainer/plugin/io.py,sha256=OmFSKLloMypletjaUr_Ptg6LS0ljqTVIp2o4Hm3eZoE,1926
@@ -149,7 +149,7 @@ nshtrainer/trainer/strategy.py,sha256=VPTn5z3zvXTydY8IJchjhjcOfpvtoejnvUkq5E4WTu
 nshtrainer/trainer/trainer.py,sha256=G_tHqzZCHJazhROcoKeOI5rZ5A8F8XlghiIWkdMbPR0,24387
 nshtrainer/util/_environment_info.py,sha256=j-wyEHKirsu3rIXTtqC2kLmIIkRe6obWjxPVWaqg2ow,24887
 nshtrainer/util/bf16.py,sha256=9QhHZCkYSfYpIcxwAMoXyuh2yTSHBzT-EdLQB297jEs,762
-nshtrainer/util/code_upload.py,sha256=CpbZEBbA8EcBElUVoCPbP5zdwtNzJhS20RLaOB-q-2k,1257
+nshtrainer/util/code_upload.py,sha256=o0GKWROL5EUvJ2F-eOr9ag6R588ZbgG8HX37fvEMfgY,1241
 nshtrainer/util/config/__init__.py,sha256=Z39JJufSb61Lhn2GfVcv3eFW_eorOrN9-9llDWlnZZM,272
 nshtrainer/util/config/dtype.py,sha256=Fn_MhhQoHPyFAnFPSwvcvLiGR3yWFIszMba02CJiC4g,2213
 nshtrainer/util/config/duration.py,sha256=mM-UfU_HvhXwW33TYEDg0x58n80tnle2e6VaWtxZTjk,764
@@ -159,6 +159,6 @@ nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
 nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
 nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
 nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
-nshtrainer-1.4.1.dist-info/METADATA,sha256=QL69Trcmw3NF3UOovpqVJbzBTtHJtnDDxAzxyj9EX24,980
-nshtrainer-1.4.1.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-nshtrainer-1.4.1.dist-info/RECORD,,
+nshtrainer-1.5.1.dist-info/METADATA,sha256=ct7S8c2O-oHJ2yw3-pApipwmP_r07z8lmFu40FhQY-k,980
+nshtrainer-1.5.1.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+nshtrainer-1.5.1.dist-info/RECORD,,