nshtrainer 1.0.0b54__py3-none-any.whl → 1.0.0b56__py3-none-any.whl

nshtrainer/configs/__init__.py

@@ -85,6 +85,7 @@ from nshtrainer.nn import NonlinearityConfig as NonlinearityConfig
 from nshtrainer.nn import NonlinearityConfigBase as NonlinearityConfigBase
 from nshtrainer.nn import PReLUConfig as PReLUConfig
 from nshtrainer.nn import ReLUNonlinearityConfig as ReLUNonlinearityConfig
+from nshtrainer.nn import RNGConfig as RNGConfig
 from nshtrainer.nn import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
 from nshtrainer.nn import SiLUNonlinearityConfig as SiLUNonlinearityConfig
 from nshtrainer.nn import SoftmaxNonlinearityConfig as SoftmaxNonlinearityConfig
@@ -306,6 +307,7 @@ __all__ = [
     "ProfilerConfig",
     "PyTorchProfilerConfig",
     "RLPSanityChecksCallbackConfig",
+    "RNGConfig",
     "ReLUNonlinearityConfig",
     "ReduceLROnPlateauConfig",
     "SLURMEnvironmentPlugin",

nshtrainer/configs/lr_scheduler/__init__.py

@@ -12,6 +12,7 @@ from nshtrainer.lr_scheduler.base import lr_scheduler_registry as lr_scheduler_r
 from nshtrainer.lr_scheduler.linear_warmup_cosine import (
     DurationConfig as DurationConfig,
 )
+from nshtrainer.lr_scheduler.reduce_lr_on_plateau import EpochsConfig as EpochsConfig
 from nshtrainer.lr_scheduler.reduce_lr_on_plateau import MetricConfig as MetricConfig
 
 from . import base as base
@@ -20,6 +21,7 @@ from . import reduce_lr_on_plateau as reduce_lr_on_plateau
 
 __all__ = [
     "DurationConfig",
+    "EpochsConfig",
     "LRSchedulerConfig",
     "LRSchedulerConfigBase",
     "LinearWarmupCosineDecayLRSchedulerConfig",

nshtrainer/configs/lr_scheduler/reduce_lr_on_plateau/__init__.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 __codegen__ = True
 
+from nshtrainer.lr_scheduler.reduce_lr_on_plateau import EpochsConfig as EpochsConfig
 from nshtrainer.lr_scheduler.reduce_lr_on_plateau import (
     LRSchedulerConfigBase as LRSchedulerConfigBase,
 )
@@ -14,6 +15,7 @@ from nshtrainer.lr_scheduler.reduce_lr_on_plateau import (
 )
 
 __all__ = [
+    "EpochsConfig",
     "LRSchedulerConfigBase",
     "MetricConfig",
     "ReduceLROnPlateauConfig",

nshtrainer/configs/nn/__init__.py

@@ -11,6 +11,7 @@ from nshtrainer.nn import NonlinearityConfig as NonlinearityConfig
 from nshtrainer.nn import NonlinearityConfigBase as NonlinearityConfigBase
 from nshtrainer.nn import PReLUConfig as PReLUConfig
 from nshtrainer.nn import ReLUNonlinearityConfig as ReLUNonlinearityConfig
+from nshtrainer.nn import RNGConfig as RNGConfig
 from nshtrainer.nn import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
 from nshtrainer.nn import SiLUNonlinearityConfig as SiLUNonlinearityConfig
 from nshtrainer.nn import SoftmaxNonlinearityConfig as SoftmaxNonlinearityConfig
@@ -25,6 +26,7 @@ from nshtrainer.nn.nonlinearity import nonlinearity_registry as nonlinearity_reg
 
 from . import mlp as mlp
 from . import nonlinearity as nonlinearity
+from . import rng as rng
 
 __all__ = [
     "ELUNonlinearityConfig",
@@ -35,6 +37,7 @@ __all__ = [
     "NonlinearityConfig",
     "NonlinearityConfigBase",
     "PReLUConfig",
+    "RNGConfig",
     "ReLUNonlinearityConfig",
     "SiLUNonlinearityConfig",
     "SigmoidNonlinearityConfig",
@@ -47,4 +50,5 @@ __all__ = [
     "mlp",
     "nonlinearity",
     "nonlinearity_registry",
+    "rng",
 ]

nshtrainer/configs/nn/rng/__init__.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.nn.rng import RNGConfig as RNGConfig
+
+__all__ = [
+    "RNGConfig",
+]

nshtrainer/configs/trainer/plugin/base/__init__.py

@@ -2,12 +2,10 @@ from __future__ import annotations
 
 __codegen__ = True
 
-from nshtrainer.trainer.plugin.base import PluginConfig as PluginConfig
 from nshtrainer.trainer.plugin.base import PluginConfigBase as PluginConfigBase
 from nshtrainer.trainer.plugin.base import plugin_registry as plugin_registry
 
 __all__ = [
-    "PluginConfig",
     "PluginConfigBase",
     "plugin_registry",
 ]

nshtrainer/configs/trainer/plugin/environment/__init__.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 __codegen__ = True
 
-from nshtrainer.trainer.plugin.environment import DTypeConfig as DTypeConfig
 from nshtrainer.trainer.plugin.environment import (
     KubeflowEnvironmentPlugin as KubeflowEnvironmentPlugin,
 )
@@ -28,7 +27,6 @@ from nshtrainer.trainer.plugin.environment import (
 from nshtrainer.trainer.plugin.environment import plugin_registry as plugin_registry
 
 __all__ = [
-    "DTypeConfig",
     "KubeflowEnvironmentPlugin",
     "LSFEnvironmentPlugin",
     "LightningEnvironmentPlugin",

nshtrainer/configs/trainer/plugin/io/__init__.py

@@ -5,7 +5,6 @@ __codegen__ = True
 from nshtrainer.trainer.plugin.io import (
     AsyncCheckpointIOPlugin as AsyncCheckpointIOPlugin,
 )
-from nshtrainer.trainer.plugin.io import PluginConfig as PluginConfig
 from nshtrainer.trainer.plugin.io import PluginConfigBase as PluginConfigBase
 from nshtrainer.trainer.plugin.io import (
     TorchCheckpointIOPlugin as TorchCheckpointIOPlugin,
@@ -15,7 +14,6 @@ from nshtrainer.trainer.plugin.io import plugin_registry as plugin_registry
 
 __all__ = [
     "AsyncCheckpointIOPlugin",
-    "PluginConfig",
     "PluginConfigBase",
     "TorchCheckpointIOPlugin",
     "XLACheckpointIOPlugin",

nshtrainer/nn/__init__.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 from .mlp import MLP as MLP
 from .mlp import MLPConfig as MLPConfig
 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
@@ -21,3 +20,5 @@ from .nonlinearity import SoftplusNonlinearityConfig as SoftplusNonlinearityConf
 from .nonlinearity import SoftsignNonlinearityConfig as SoftsignNonlinearityConfig
 from .nonlinearity import SwishNonlinearityConfig as SwishNonlinearityConfig
 from .nonlinearity import TanhNonlinearityConfig as TanhNonlinearityConfig
+from .rng import RNGConfig as RNGConfig
+from .rng import rng_context as rng_context

nshtrainer/nn/rng.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import contextlib
+
+import nshconfig as C
+import torch
+
+
+@contextlib.contextmanager
+def rng_context(config: RNGConfig | None):
+    with contextlib.ExitStack() as stack:
+        if config is not None:
+            stack.enter_context(
+                torch.random.fork_rng(devices=range(torch.cuda.device_count()))
+            )
+            torch.manual_seed(config.seed)
+
+        yield
+
+
+class RNGConfig(C.Config):
+    seed: int
+    """Random seed to use for initialization."""

nshtrainer/trainer/plugin/__init__.py

@@ -1,10 +1,18 @@
 from __future__ import annotations
 
+from typing import Annotated
+
+from typing_extensions import TypeAliasType
+
 from . import environment as environment
 from . import io as io
 from . import layer_sync as layer_sync
 from . import precision as precision
 from .base import Plugin as Plugin
-from .base import PluginConfig as PluginConfig
 from .base import PluginConfigBase as PluginConfigBase
 from .base import plugin_registry as plugin_registry
+
+PluginConfig = TypeAliasType(
+    "PluginConfig",
+    Annotated[PluginConfigBase, plugin_registry.DynamicResolution()],
+)

nshtrainer/trainer/plugin/base.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-import logging
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Annotated
+from typing import TYPE_CHECKING
 
 import nshconfig as C
 from lightning.fabric.plugins import CheckpointIO, ClusterEnvironment
@@ -12,7 +11,6 @@ from typing_extensions import TypeAliasType
 
 if TYPE_CHECKING:
     from .._config import TrainerConfig
-log = logging.getLogger(__name__)
 
 
 Plugin = TypeAliasType(
@@ -26,8 +24,3 @@ class PluginConfigBase(C.Config, ABC):
 
 
 plugin_registry = C.Registry(PluginConfigBase, discriminator="name")
-
-PluginConfig = TypeAliasType(
-    "PluginConfig",
-    Annotated[PluginConfigBase, plugin_registry.DynamicResolution()],
-)

nshtrainer/trainer/plugin/environment.py

@@ -1,27 +1,25 @@
 from __future__ import annotations
 
 import signal
-from typing import Any, Literal
+from typing import Literal
 
-from lightning.pytorch.plugins.environments import ClusterEnvironment
-from typing_extensions import override
+from typing_extensions import TypeAliasType, override
 
-from ...util.config.dtype import DTypeConfig
 from .base import PluginConfigBase, plugin_registry
 
 
 @plugin_registry.register
 class KubeflowEnvironmentPlugin(PluginConfigBase):
-    name: Literal["kubeflow_environment"] = "kubeflow_environment"
-
     """Environment for distributed training using the PyTorchJob operator from Kubeflow.
 
     This environment, unlike others, does not get auto-detected and needs to be passed
     to the Fabric/Trainer constructor manually.
     """
 
+    name: Literal["kubeflow_environment"] = "kubeflow_environment"
+
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.kubeflow import KubeflowEnvironment
 
         return KubeflowEnvironment()
@@ -29,8 +27,6 @@ class KubeflowEnvironmentPlugin(PluginConfigBase):
 
 @plugin_registry.register
 class LightningEnvironmentPlugin(PluginConfigBase):
-    name: Literal["lightning_environment"] = "lightning_environment"
-
     """The default environment used by Lightning for a single node or free cluster (not managed).
 
     There are two modes the Lightning environment can operate with:
@@ -40,8 +36,10 @@ class LightningEnvironmentPlugin(PluginConfigBase):
        The appropriate environment variables need to be set, and at minimum `LOCAL_RANK`.
     """
 
+    name: Literal["lightning_environment"] = "lightning_environment"
+
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.lightning import LightningEnvironment
 
         return LightningEnvironment()
@@ -49,16 +47,16 @@ class LightningEnvironmentPlugin(PluginConfigBase):
 
 @plugin_registry.register
 class LSFEnvironmentPlugin(PluginConfigBase):
-    name: Literal["lsf_environment"] = "lsf_environment"
-
     """An environment for running on clusters managed by the LSF resource manager.
 
     It is expected that any execution using this ClusterEnvironment was executed
     using the Job Step Manager i.e. `jsrun`.
     """
 
+    name: Literal["lsf_environment"] = "lsf_environment"
+
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.lsf import LSFEnvironment
 
         return LSFEnvironment()
@@ -66,48 +64,108 @@ class LSFEnvironmentPlugin(PluginConfigBase):
 
 @plugin_registry.register
 class MPIEnvironmentPlugin(PluginConfigBase):
-    name: Literal["mpi_environment"] = "mpi_environment"
-
     """An environment for running on clusters with processes created through MPI.
 
     Requires the installation of the `mpi4py` package.
     """
 
+    name: Literal["mpi_environment"] = "mpi_environment"
+
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.mpi import MPIEnvironment
 
         return MPIEnvironment()
 
 
+SignalAlias = TypeAliasType(
+    "SignalAlias",
+    Literal[
+        "SIGABRT",
+        "SIGFPE",
+        "SIGILL",
+        "SIGINT",
+        "SIGSEGV",
+        "SIGTERM",
+        "SIGBREAK",
+        "CTRL_C_EVENT",
+        "CTRL_BREAK_EVENT",
+        "SIGALRM",
+        "SIGBUS",
+        "SIGCHLD",
+        "SIGCONT",
+        "SIGHUP",
+        "SIGIO",
+        "SIGIOT",
+        "SIGKILL",
+        "SIGPIPE",
+        "SIGPROF",
+        "SIGQUIT",
+        "SIGSTOP",
+        "SIGSYS",
+        "SIGTRAP",
+        "SIGTSTP",
+        "SIGTTIN",
+        "SIGTTOU",
+        "SIGURG",
+        "SIGUSR1",
+        "SIGUSR2",
+        "SIGVTALRM",
+        "SIGWINCH",
+        "SIGXCPU",
+        "SIGXFSZ",
+        "SIGEMT",
+        "SIGINFO",
+        "SIGCLD",
+        "SIGPOLL",
+        "SIGPWR",
+        "SIGRTMAX",
+        "SIGRTMIN",
+        "SIGSTKFLT",
+    ],
+)
+
+
 @plugin_registry.register
 class SLURMEnvironmentPlugin(PluginConfigBase):
+    """An environment for running on clusters managed by the SLURM resource manager."""
+
     name: Literal["slurm_environment"] = "slurm_environment"
 
     auto_requeue: bool = True
     """Whether automatic job resubmission is enabled or not."""
 
-    requeue_signal: signal.Signals | None = None
+    requeue_signal: SignalAlias | None = None
     """The signal that SLURM will send to indicate that the job should be requeued."""
 
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
 
+        requeue_signal = None
+        if self.requeue_signal is not None:
+            try:
+                requeue_signal = signal.Signals[self.requeue_signal]
+            except KeyError:
+                raise ValueError(
+                    f"Invalid signal name: {self.requeue_signal}. "
+                    "Please provide a valid signal name from the signal module."
+                )
+
         return SLURMEnvironment(
             auto_requeue=self.auto_requeue,
-            requeue_signal=self.requeue_signal,
+            requeue_signal=requeue_signal,
         )
 
 
 @plugin_registry.register
 class TorchElasticEnvironmentPlugin(PluginConfigBase):
-    name: Literal["torchelastic_environment"] = "torchelastic_environment"
-
     """Environment for fault-tolerant and elastic training with torchelastic."""
 
+    name: Literal["torchelastic_environment"] = "torchelastic_environment"
+
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.torchelastic import (
             TorchElasticEnvironment,
         )
@@ -117,12 +175,12 @@ class TorchElasticEnvironmentPlugin(PluginConfigBase):
 
 @plugin_registry.register
 class XLAEnvironmentPlugin(PluginConfigBase):
-    name: Literal["xla_environment"] = "xla_environment"
-
     """Cluster environment for training on a TPU Pod with the PyTorch/XLA library."""
 
+    name: Literal["xla_environment"] = "xla_environment"
+
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.xla import XLAEnvironment
 
         return XLAEnvironment()

nshtrainer/trainer/plugin/io.py

@@ -2,27 +2,52 @@ from __future__ import annotations
 
 from typing import Literal
 
-from lightning.pytorch.plugins.io import CheckpointIO
 from typing_extensions import override
 
-from .base import PluginConfig, PluginConfigBase, plugin_registry
+from .base import PluginConfigBase, plugin_registry
 
 
 @plugin_registry.register
-class AsyncCheckpointIOPlugin(PluginConfigBase):
-    name: Literal["async_checkpoint"] = "async_checkpoint"
+class TorchCheckpointIOPlugin(PluginConfigBase):
+    """CheckpointIO that utilizes torch.save and torch.load to save and load checkpoints respectively."""
+
+    name: Literal["torch_checkpoint"] = "torch_checkpoint"
+
+    @override
+    def create_plugin(self, trainer_config):
+        from lightning.fabric.plugins.io.torch_io import TorchCheckpointIO
+
+        return TorchCheckpointIO()
+
+
+@plugin_registry.register
+class XLACheckpointIOPlugin(PluginConfigBase):
+    """CheckpointIO that utilizes xm.save to save checkpoints for TPU training strategies."""
 
+    name: Literal["xla_checkpoint"] = "xla_checkpoint"
+
+    @override
+    def create_plugin(self, trainer_config):
+        from lightning.fabric.plugins.io.xla import XLACheckpointIO
+
+        return XLACheckpointIO()
+
+
+@plugin_registry.register
+class AsyncCheckpointIOPlugin(PluginConfigBase):
     """Enables saving the checkpoints asynchronously in a thread.
 
     .. warning::  This is an experimental feature.
     """
 
-    checkpoint_io: PluginConfig | None = None
+    name: Literal["async_checkpoint"] = "async_checkpoint"
+
+    checkpoint_io: TorchCheckpointIOPlugin | None = None
     """A checkpoint IO plugin that is used as the basis for async checkpointing."""
 
     @override
-    def create_plugin(self, trainer_config) -> CheckpointIO:
-        from lightning.pytorch.plugins.io.async_plugin import AsyncCheckpointIO
+    def create_plugin(self, trainer_config):
+        from lightning.pytorch.plugins.io import AsyncCheckpointIO, CheckpointIO
 
         base_io = (
             self.checkpoint_io.create_plugin(trainer_config)
@@ -34,29 +59,3 @@ class AsyncCheckpointIOPlugin(PluginConfigBase):
                 f"Expected `checkpoint_io` to be a `CheckpointIO` instance, but got {type(base_io)}."
             )
         return AsyncCheckpointIO(checkpoint_io=base_io)
-
-
-@plugin_registry.register
-class TorchCheckpointIOPlugin(PluginConfigBase):
-    name: Literal["torch_checkpoint"] = "torch_checkpoint"
-
-    """CheckpointIO that utilizes torch.save and torch.load to save and load checkpoints respectively."""
-
-    @override
-    def create_plugin(self, trainer_config) -> CheckpointIO:
-        from lightning.fabric.plugins.io.torch_io import TorchCheckpointIO
-
-        return TorchCheckpointIO()
-
-
-@plugin_registry.register
-class XLACheckpointIOPlugin(PluginConfigBase):
-    name: Literal["xla_checkpoint"] = "xla_checkpoint"
-
-    """CheckpointIO that utilizes xm.save to save checkpoints for TPU training strategies."""
-
-    @override
-    def create_plugin(self, trainer_config) -> CheckpointIO:
-        from lightning.fabric.plugins.io.xla import XLACheckpointIO
-
-        return XLACheckpointIO()

nshtrainer/trainer/plugin/layer_sync.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 from typing import Literal
 
-from lightning.pytorch.plugins.layer_sync import LayerSync
 from typing_extensions import override
 
 from .base import PluginConfigBase, plugin_registry
@@ -10,16 +9,16 @@ from .base import PluginConfigBase, plugin_registry
 
 @plugin_registry.register
 class TorchSyncBatchNormPlugin(PluginConfigBase):
-    name: Literal["torch_sync_batchnorm"] = "torch_sync_batchnorm"
-
     """A plugin that wraps all batch normalization layers of a model with synchronization
     logic for multiprocessing.
 
     This plugin has no effect in single-device operation.
     """
 
+    name: Literal["torch_sync_batchnorm"] = "torch_sync_batchnorm"
+
     @override
-    def create_plugin(self, trainer_config) -> LayerSync:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.layer_sync import TorchSyncBatchNorm
 
         return TorchSyncBatchNorm()

nshtrainer/trainer/plugin/precision.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 from typing import Any, Literal
 
-from lightning.pytorch.plugins.precision import Precision
 from typing_extensions import override
 
 from ...util.config.dtype import DTypeConfig
@@ -20,7 +19,7 @@ class MixedPrecisionPluginConfig(PluginConfigBase):
     """The device for ``torch.autocast``."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.amp import MixedPrecision
 
         return MixedPrecision(self.precision, self.device)
@@ -45,7 +44,7 @@ class BitsandbytesPluginConfig(PluginConfigBase):
     """
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.bitsandbytes import (
             BitsandbytesPrecision,
         )
@@ -66,7 +65,7 @@ class DeepSpeedPluginConfig(PluginConfigBase):
     mixed precision (16-mixed, bf16-mixed)."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.deepspeed import DeepSpeedPrecision
 
         return DeepSpeedPrecision(precision=self.precision)
@@ -80,7 +79,7 @@ class DoublePrecisionPluginConfig(PluginConfigBase):
     """Plugin for training with double (``torch.float64``) precision."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.double import DoublePrecision
 
         return DoublePrecision()
@@ -95,7 +94,7 @@ class FSDPPrecisionPluginConfig(PluginConfigBase):
     mixed precision (16-mixed, bf16-mixed)."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.fsdp import FSDPPrecision
 
         return FSDPPrecision(precision=self.precision)
@@ -109,7 +108,7 @@ class HalfPrecisionPluginConfig(PluginConfigBase):
     """Whether to use ``torch.float16`` (``'16-true'``) or ``torch.bfloat16`` (``'bf16-true'``)."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.half import HalfPrecision
 
         return HalfPrecision(precision=self.precision)
@@ -134,7 +133,7 @@ class TransformerEnginePluginConfig(PluginConfigBase):
     Defaults to the same as weights_dtype."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.transformer_engine import (
             TransformerEnginePrecision,
         )
@@ -157,7 +156,7 @@ class XLAPluginConfig(PluginConfigBase):
     """Full precision (32-true) or half precision (16-true, bf16-true)."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.xla import XLAPrecision
 
         return XLAPrecision(precision=self.precision)

{nshtrainer-1.0.0b54.dist-info → nshtrainer-1.0.0b56.dist-info}/METADATA

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

{nshtrainer-1.0.0b54.dist-info → nshtrainer-1.0.0b56.dist-info}/RECORD

@@ -32,7 +32,7 @@ nshtrainer/callbacks/timer.py,sha256=gDcw_K_ikf0bkVgxQ0cDhvvNvz6GLZVLcatuKfh0ORU
 nshtrainer/callbacks/wandb_upload_code.py,sha256=shV7UtnXgY2bUlXdVrXiaDs0PNLlIt7TzNJkJPkzvzI,2414
 nshtrainer/callbacks/wandb_watch.py,sha256=VB14Dy5ZRXQ3di0fPv0K_DFJurLhroLPytnuwQBiJFg,3037
 nshtrainer/configs/.gitattributes,sha256=VeZmarvNEqiRBOHGcllpKm90nL6C8u4tBu7SEm7fj-E,26
-nshtrainer/configs/__init__.py,sha256=0BzCgE1iEJ0Ywmy__mqJZipLQtwZVdz6XK-gHbkA7GY,14650
+nshtrainer/configs/__init__.py,sha256=-rGk9pnRnuz4yKvACGOpY3nkrWnHholqZGk7UP2Vkrc,14716
 nshtrainer/configs/_checkpoint/__init__.py,sha256=6s7Y68StboqscY2G4P_QG443jz5aiym5SjOogIljWLg,342
 nshtrainer/configs/_checkpoint/metadata/__init__.py,sha256=oOPfYkXTjKgm6pluGsG6V1TPyCEGjsQpHVL-LffSUFQ,290
 nshtrainer/configs/_directory/__init__.py,sha256=_oO7vM9DhzHSxtZcv86sTi7hZIptnK1gr-AP9mqQ370,386
@@ -67,15 +67,16 @@ nshtrainer/configs/loggers/base/__init__.py,sha256=HLUfEDbjaAXqzsFmQbjdciIWzR1st
 nshtrainer/configs/loggers/csv/__init__.py,sha256=gawaDX92JObGSmBqYpfNHWMHBwVOofS694W-1Y2GWDU,353
 nshtrainer/configs/loggers/tensorboard/__init__.py,sha256=phzm-TnBkdkibTgoOxIIcAliqL3zU8gSNK61Mwxs1CM,410
 nshtrainer/configs/loggers/wandb/__init__.py,sha256=TDcD5WZSKenc2mgIXhwz2l96l8P_Ur3N5CzEol5AKGw,746
-nshtrainer/configs/lr_scheduler/__init__.py,sha256=xtiUx0isxA82-uXMn4-KmPnDCfbUkpAnd2_pFupAAKQ,1137
+nshtrainer/configs/lr_scheduler/__init__.py,sha256=PvH2d8QEC3TsC3_svcUbxeQEMMzIK_In0_Bp9xntSms,1243
 nshtrainer/configs/lr_scheduler/base/__init__.py,sha256=6Cx8r4rdxeSYxc_z0o7drKCblGJU_zzqrOoYlWYR5qY,305
 nshtrainer/configs/lr_scheduler/linear_warmup_cosine/__init__.py,sha256=5ZMLDO9VL6SNU6pF-62lDnpmqix3_Ol9DdEwiuOPYlA,675
-nshtrainer/configs/lr_scheduler/reduce_lr_on_plateau/__init__.py,sha256=w-vq8UbRGPX8DZVWCMC5eIrbvVc_guxjj7Du9AaeKCw,609
+nshtrainer/configs/lr_scheduler/reduce_lr_on_plateau/__init__.py,sha256=DBwZV590I5qwyOS5M43YhUzgYy1-AjzkM5aEnTA6XdI,715
 nshtrainer/configs/metrics/__init__.py,sha256=mK_xgXJDAyGY4K_x_Zo4aj36kjT45d850keuUe3U1rY,200
 nshtrainer/configs/metrics/_config/__init__.py,sha256=XDDvDPWULd_vd3lrgF2KGAVR2LVDuhdQvy-fF2ImarI,159
-nshtrainer/configs/nn/__init__.py,sha256=tkFG2Hb0oL_AmWP3_0WkDN2zI5PkVfrgwXhaAII7CZw,2072
+nshtrainer/configs/nn/__init__.py,sha256=Ms2gIqbRxNVm6GHKCddCJTTqMwUPifjjHD_fCfJpcMQ,2174
 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=itIDIHQvGm50eZ7JLyNElahnNUMPJ__4PMmTjc0RQ6o,444
 nshtrainer/configs/profiler/__init__.py,sha256=2ssaIpfVnvcbfNvZ-JeKp1Cx4NO1LknkVqTm1hu7Lvw,768
 nshtrainer/configs/profiler/_base/__init__.py,sha256=ekYfPg-VDhCAFM5nJka2TxUYdRDm1CKqjwUOQNbQjD4,176
@@ -86,9 +87,9 @@ nshtrainer/configs/trainer/__init__.py,sha256=a8pzGVid52abAVARPbgjaN566H1ZM44FH_
 nshtrainer/configs/trainer/_config/__init__.py,sha256=6DXdtP-uH11TopQ7kzId9fco-wVkD7ZfevbBqDpN6TE,3817
 nshtrainer/configs/trainer/accelerator/__init__.py,sha256=3H6R3wlwbKL1TzDqGCChZk78-BcE2czLouo7Djiq3nA,898
 nshtrainer/configs/trainer/plugin/__init__.py,sha256=NkHQxMPkrtTtdIAO4dQUE9SWEcHRDB0yUXLkTjnl4dA,3332
-nshtrainer/configs/trainer/plugin/base/__init__.py,sha256=GuubcKrbXt4VjJRT8VpNUQqBtyuutre_CfkS0EWZ5_E,368
-nshtrainer/configs/trainer/plugin/environment/__init__.py,sha256=3o16x4qRAOvkJH9Vg4-QwsEODDC6aP_OXRnPPkm_xSo,1376
-nshtrainer/configs/trainer/plugin/io/__init__.py,sha256=W6G67JnigB6d3MiwLrbSKgtIZLUccXznp-IXwkK1J4U,743
+nshtrainer/configs/trainer/plugin/base/__init__.py,sha256=slW5z1FZw2qICXO9l9DnLIDB1Yl7KOcxPEZkyYIHrp4,276
+nshtrainer/configs/trainer/plugin/environment/__init__.py,sha256=geTaOfbFep3DVDgSsRrufM7q6MIZur5QG1_47ngAL0I,1280
+nshtrainer/configs/trainer/plugin/io/__init__.py,sha256=AtGUuE0M16dTpX0q9NqvJiE4qU1j07N0RLtd-JFzWuc,653
 nshtrainer/configs/trainer/plugin/layer_sync/__init__.py,sha256=SYDZk2M6sgpt4sEuoURuS8EKYmaqGcvYxETE9jvTrEE,431
 nshtrainer/configs/trainer/plugin/precision/__init__.py,sha256=szlqSfK2XuWdkf72LQzQFv3SlWfKFdRUpBEYIxQ3TPs,1507
 nshtrainer/configs/trainer/strategy/__init__.py,sha256=50whNloJVBq_bdbLaPQnPBTeS1Rcs8MwxTCYBj1kKa4,273
@@ -119,11 +120,12 @@ nshtrainer/model/base.py,sha256=bZMNap0rkxRbAbu2BOHV_6YS2iZZnvy6wVSMOXGa_ZM,8680
 nshtrainer/model/mixins/callback.py,sha256=0LPgve4VszHbLipid4mpI1qnnmdGS2spivs0dXLvqHw,3154
 nshtrainer/model/mixins/debug.py,sha256=ydLuAAaa7M5bX0gougZ5gWuZnvn4Ra9assal3IZ9hq8,2086
 nshtrainer/model/mixins/logger.py,sha256=7u9fQig-SVFA9RFIB4U0gqJAzruh49mgmXXvZ6VkDUk,11694
-nshtrainer/nn/__init__.py,sha256=5Gg3nieGSC5_dXaI9KUVUUbM13hHexH9831m4hcf6no,1475
+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/rng.py,sha256=IJGvX9v8qBkfgBrMlNU2aj-MbYTPoncFyJzvPkzCQpM,512
 nshtrainer/optimizer.py,sha256=u968GRNPUNn3f_9BEY2RBNuJq5O3wJWams3NG0dkrOA,1738
 nshtrainer/profiler/__init__.py,sha256=RjaNBoVcTFu8lF0dNlFp-2LaPYdonoIbDy2_KhgF0Ek,594
 nshtrainer/profiler/_base.py,sha256=kFcSVn9gJuMwgDxbfyHh46CmEAIPZjxw3yjPbKgzvwA,950
@@ -134,12 +136,12 @@ nshtrainer/trainer/__init__.py,sha256=fQ7gQRlGWX-90TYT0rttkQyvXDCzo7DAvJgr-jX1zs
 nshtrainer/trainer/_config.py,sha256=s-_XoLc9mbNAdroRJyOKd3dLTyrFLQkPyGJkKDmBYf8,33267
 nshtrainer/trainer/_runtime_callback.py,sha256=6F2Gq27Q8OFfN3RtdNC6QRA8ac0LC1hh4DUE3V5WgbI,4217
 nshtrainer/trainer/accelerator.py,sha256=Bqq-ry7DeCY4zw9_zBvTZiijpA-uUHrDjtbLV652m4M,2415
-nshtrainer/trainer/plugin/__init__.py,sha256=UM8f70Ml3RGpsXeQr1Yh1yBcxxjFNGFGgJbniCn_rws,366
-nshtrainer/trainer/plugin/base.py,sha256=9-qUHXGpll_yCylun0899sbmJDpyhD9IQcBtVrJx38I,919
-nshtrainer/trainer/plugin/environment.py,sha256=NW0qbsbvDPe59JGOMgPLq1fj7szLucIV1WRTxCrcjF4,4367
-nshtrainer/trainer/plugin/io.py,sha256=nm6YDCVZAhmPvLaLnw6q4BrK2Gj2wvD5ZLDhj1xneEE,2030
-nshtrainer/trainer/plugin/layer_sync.py,sha256=h-ydZwXepnsw5-paLgiDatqPyQ_8C0QEv6DrYFIaXOo,735
-nshtrainer/trainer/plugin/precision.py,sha256=I0QsB1bVxmsFmBOkgrAfGONsuYae_lD9Bz0PfJEQvH4,5598
+nshtrainer/trainer/plugin/__init__.py,sha256=q_q98MYNaZ2VE_tqGqYlQjQnlaF4NE1FUqVVbj0EK7k,517
+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
+nshtrainer/trainer/plugin/layer_sync.py,sha256=-BbEyWZ063O7tZme7Gdu1lVxK6p1NeuLcOfPXnvH29s,663
+nshtrainer/trainer/plugin/precision.py,sha256=7lf7KZd_yFyPmhLApjEIv0pkoDB5zdxi-7in0wRj3z8,5436
 nshtrainer/trainer/signal_connector.py,sha256=GhfGcSzfaTNhnj2QFkBDq5aT7FqbLMA7eC8SYQs8_8w,10828
 nshtrainer/trainer/strategy.py,sha256=VPTn5z3zvXTydY8IJchjhjcOfpvtoejnvUkq5E4WTus,1368
 nshtrainer/trainer/trainer.py,sha256=Lo3vUo3ooTAjaX2fUYPFSMv5FP7sWfVov0QbA-T5hZ8,21113
@@ -154,6 +156,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.0.0b54.dist-info/METADATA,sha256=oSfrN2tgKgkZJwGbZVNbLULQcVMxh_wb02u7Hrujfn4,988
-nshtrainer-1.0.0b54.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-nshtrainer-1.0.0b54.dist-info/RECORD,,
+nshtrainer-1.0.0b56.dist-info/METADATA,sha256=KevxOAySDQL3f1OO7dgEsNUeX322ostuFo6i7egHzRU,988
+nshtrainer-1.0.0b56.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+nshtrainer-1.0.0b56.dist-info/RECORD,,