nshtrainer 1.0.0b26__py3-none-any.whl → 1.0.0b28__py3-none-any.whl

nshtrainer/callbacks/actsave.py

@@ -4,12 +4,12 @@ import contextlib
 from pathlib import Path
 from typing import Literal
 
-from typing_extensions import TypeAlias, override
+from typing_extensions import TypeAliasType, override
 
 from .._callback import NTCallbackBase
 from .base import CallbackConfigBase
 
-Stage: TypeAlias = Literal["train", "validation", "test", "predict"]
+Stage = TypeAliasType("Stage", Literal["train", "validation", "test", "predict"])
 
 
 class ActSaveConfig(CallbackConfigBase):

nshtrainer/callbacks/base.py

@@ -4,11 +4,11 @@ from abc import ABC, abstractmethod
 from collections import Counter
 from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, ClassVar, TypeAlias
+from typing import TYPE_CHECKING, ClassVar
 
 import nshconfig as C
 from lightning.pytorch import Callback
-from typing_extensions import TypedDict, Unpack
+from typing_extensions import TypeAliasType, TypedDict, Unpack
 
 if TYPE_CHECKING:
     from ..trainer._config import TrainerConfig
@@ -30,7 +30,9 @@ class CallbackWithMetadata:
     metadata: CallbackMetadataConfig
 
 
-ConstructedCallback: TypeAlias = Callback | CallbackWithMetadata
+ConstructedCallback = TypeAliasType(
+    "ConstructedCallback", Callback | CallbackWithMetadata
+)
 
 
 class CallbackConfigBase(C.Config, ABC):

nshtrainer/callbacks/shared_parameters.py

@@ -2,12 +2,12 @@ from __future__ import annotations
 
 import logging
 from collections.abc import Iterable
-from typing import Literal, Protocol, TypeAlias, runtime_checkable
+from typing import Literal, Protocol, runtime_checkable
 
 import torch.nn as nn
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import Callback
-from typing_extensions import override
+from typing_extensions import TypeAliasType, override
 
 from .base import CallbackConfigBase
 
@@ -34,7 +34,9 @@ class SharedParametersCallbackConfig(CallbackConfigBase):
         yield SharedParametersCallback(self)
 
 
-SharedParametersList: TypeAlias = list[tuple[nn.Parameter, int | float]]
+SharedParametersList = TypeAliasType(
+    "SharedParametersList", list[tuple[nn.Parameter, int | float]]
+)
 
 
 @runtime_checkable

nshtrainer/configs/__init__.py

@@ -14,6 +14,7 @@ from nshtrainer._hf_hub import HuggingFaceHubConfig as HuggingFaceHubConfig
 from nshtrainer.callbacks import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
+from nshtrainer.callbacks import CallbackConfig as CallbackConfig
 from nshtrainer.callbacks import DebugFlagCallbackConfig as DebugFlagCallbackConfig
 from nshtrainer.callbacks import (
     DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
@@ -46,6 +47,9 @@ from nshtrainer.callbacks import (
 from nshtrainer.callbacks import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
+from nshtrainer.callbacks import (
+    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
+)
 from nshtrainer.callbacks import (
     WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
 )
@@ -57,11 +61,13 @@ from nshtrainer.callbacks.checkpoint._base import (
 from nshtrainer.loggers import ActSaveLoggerConfig as ActSaveLoggerConfig
 from nshtrainer.loggers import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.loggers import CSVLoggerConfig as CSVLoggerConfig
+from nshtrainer.loggers import LoggerConfig as LoggerConfig
 from nshtrainer.loggers import TensorboardLoggerConfig as TensorboardLoggerConfig
 from nshtrainer.loggers import WandbLoggerConfig as WandbLoggerConfig
 from nshtrainer.lr_scheduler import (
     LinearWarmupCosineDecayLRSchedulerConfig as LinearWarmupCosineDecayLRSchedulerConfig,
 )
+from nshtrainer.lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
 from nshtrainer.lr_scheduler import LRSchedulerConfigBase as LRSchedulerConfigBase
 from nshtrainer.lr_scheduler import ReduceLROnPlateauConfig as ReduceLROnPlateauConfig
 from nshtrainer.nn import BaseNonlinearityConfig as BaseNonlinearityConfig
@@ -70,6 +76,7 @@ from nshtrainer.nn import GELUNonlinearityConfig as GELUNonlinearityConfig
 from nshtrainer.nn import LeakyReLUNonlinearityConfig as LeakyReLUNonlinearityConfig
 from nshtrainer.nn import MishNonlinearityConfig as MishNonlinearityConfig
 from nshtrainer.nn import MLPConfig as MLPConfig
+from nshtrainer.nn import NonlinearityConfig as NonlinearityConfig
 from nshtrainer.nn import PReLUConfig as PReLUConfig
 from nshtrainer.nn import ReLUNonlinearityConfig as ReLUNonlinearityConfig
 from nshtrainer.nn import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
@@ -83,12 +90,17 @@ from nshtrainer.nn.nonlinearity import (
     SwiGLUNonlinearityConfig as SwiGLUNonlinearityConfig,
 )
 from nshtrainer.optimizer import AdamWConfig as AdamWConfig
+from nshtrainer.optimizer import OptimizerConfig as OptimizerConfig
 from nshtrainer.optimizer import OptimizerConfigBase as OptimizerConfigBase
 from nshtrainer.profiler import AdvancedProfilerConfig as AdvancedProfilerConfig
 from nshtrainer.profiler import BaseProfilerConfig as BaseProfilerConfig
+from nshtrainer.profiler import ProfilerConfig as ProfilerConfig
 from nshtrainer.profiler import PyTorchProfilerConfig as PyTorchProfilerConfig
 from nshtrainer.profiler import SimpleProfilerConfig as SimpleProfilerConfig
 from nshtrainer.trainer._config import AcceleratorConfigBase as AcceleratorConfigBase
+from nshtrainer.trainer._config import (
+    CheckpointCallbackConfig as CheckpointCallbackConfig,
+)
 from nshtrainer.trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
 from nshtrainer.trainer._config import EnvironmentConfig as EnvironmentConfig
 from nshtrainer.trainer._config import GradientClippingConfig as GradientClippingConfig
@@ -127,6 +139,7 @@ from nshtrainer.util._environment_info import (
 )
 from nshtrainer.util._environment_info import GitRepositoryConfig as GitRepositoryConfig
 from nshtrainer.util.config import DTypeConfig as DTypeConfig
+from nshtrainer.util.config import DurationConfig as DurationConfig
 from nshtrainer.util.config import EpochsConfig as EpochsConfig
 from nshtrainer.util.config import StepsConfig as StepsConfig
 
@@ -155,13 +168,16 @@ __all__ = [
     "BaseProfilerConfig",
     "BestCheckpointCallbackConfig",
     "CSVLoggerConfig",
+    "CallbackConfig",
     "CallbackConfigBase",
+    "CheckpointCallbackConfig",
     "CheckpointMetadata",
     "CheckpointSavingConfig",
     "DTypeConfig",
     "DebugFlagCallbackConfig",
     "DirectoryConfig",
     "DirectorySetupCallbackConfig",
+    "DurationConfig",
     "ELUNonlinearityConfig",
     "EMACallbackConfig",
     "EarlyStoppingCallbackConfig",
@@ -184,21 +200,26 @@ __all__ = [
     "GradientSkippingCallbackConfig",
     "HuggingFaceHubAutoCreateConfig",
     "HuggingFaceHubConfig",
+    "LRSchedulerConfig",
     "LRSchedulerConfigBase",
     "LastCheckpointCallbackConfig",
     "LeakyReLUNonlinearityConfig",
     "LearningRateMonitorConfig",
     "LinearWarmupCosineDecayLRSchedulerConfig",
     "LogEpochCallbackConfig",
+    "LoggerConfig",
     "MLPConfig",
     "MetricConfig",
     "MishNonlinearityConfig",
+    "NonlinearityConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
+    "OptimizerConfig",
     "OptimizerConfigBase",
     "PReLUConfig",
     "PluginConfigBase",
     "PrintTableMetricsCallbackConfig",
+    "ProfilerConfig",
     "PyTorchProfilerConfig",
     "RLPSanityChecksCallbackConfig",
     "ReLUNonlinearityConfig",
@@ -217,6 +238,7 @@ __all__ = [
     "SwishNonlinearityConfig",
     "TanhNonlinearityConfig",
     "TensorboardLoggerConfig",
+    "TimeCheckpointCallbackConfig",
     "TrainerConfig",
     "WandbLoggerConfig",
     "WandbUploadCodeCallbackConfig",

nshtrainer/configs/_directory/__init__.py

@@ -6,8 +6,10 @@ from nshtrainer._directory import DirectoryConfig as DirectoryConfig
 from nshtrainer._directory import (
     DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
 )
+from nshtrainer._directory import LoggerConfig as LoggerConfig
 
 __all__ = [
     "DirectoryConfig",
     "DirectorySetupCallbackConfig",
+    "LoggerConfig",
 ]

nshtrainer/configs/callbacks/__init__.py

@@ -5,6 +5,7 @@ __codegen__ = True
 from nshtrainer.callbacks import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
+from nshtrainer.callbacks import CallbackConfig as CallbackConfig
 from nshtrainer.callbacks import CallbackConfigBase as CallbackConfigBase
 from nshtrainer.callbacks import DebugFlagCallbackConfig as DebugFlagCallbackConfig
 from nshtrainer.callbacks import (
@@ -38,6 +39,9 @@ from nshtrainer.callbacks import (
 from nshtrainer.callbacks import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
+from nshtrainer.callbacks import (
+    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
+)
 from nshtrainer.callbacks import (
     WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
 )
@@ -77,6 +81,7 @@ __all__ = [
     "ActSaveConfig",
     "BaseCheckpointCallbackConfig",
     "BestCheckpointCallbackConfig",
+    "CallbackConfig",
     "CallbackConfigBase",
     "CheckpointMetadata",
     "DebugFlagCallbackConfig",
@@ -95,6 +100,7 @@ __all__ = [
     "PrintTableMetricsCallbackConfig",
     "RLPSanityChecksCallbackConfig",
     "SharedParametersCallbackConfig",
+    "TimeCheckpointCallbackConfig",
     "WandbUploadCodeCallbackConfig",
     "WandbWatchCallbackConfig",
     "actsave",

nshtrainer/configs/callbacks/checkpoint/__init__.py

@@ -11,6 +11,9 @@ from nshtrainer.callbacks.checkpoint import (
 from nshtrainer.callbacks.checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
+from nshtrainer.callbacks.checkpoint import (
+    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
+)
 from nshtrainer.callbacks.checkpoint._base import (
     BaseCheckpointCallbackConfig as BaseCheckpointCallbackConfig,
 )
@@ -26,6 +29,7 @@ from . import _base as _base
 from . import best_checkpoint as best_checkpoint
 from . import last_checkpoint as last_checkpoint
 from . import on_exception_checkpoint as on_exception_checkpoint
+from . import time_checkpoint as time_checkpoint
 
 __all__ = [
     "BaseCheckpointCallbackConfig",
@@ -35,8 +39,10 @@ __all__ = [
     "LastCheckpointCallbackConfig",
     "MetricConfig",
     "OnExceptionCheckpointCallbackConfig",
+    "TimeCheckpointCallbackConfig",
     "_base",
     "best_checkpoint",
     "last_checkpoint",
     "on_exception_checkpoint",
+    "time_checkpoint",
 ]

nshtrainer/configs/callbacks/checkpoint/time_checkpoint/__init__.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.callbacks.checkpoint.time_checkpoint import (
+    BaseCheckpointCallbackConfig as BaseCheckpointCallbackConfig,
+)
+from nshtrainer.callbacks.checkpoint.time_checkpoint import (
+    CheckpointMetadata as CheckpointMetadata,
+)
+from nshtrainer.callbacks.checkpoint.time_checkpoint import (
+    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
+)
+
+__all__ = [
+    "BaseCheckpointCallbackConfig",
+    "CheckpointMetadata",
+    "TimeCheckpointCallbackConfig",
+]

nshtrainer/configs/loggers/__init__.py

@@ -5,6 +5,7 @@ __codegen__ = True
 from nshtrainer.loggers import ActSaveLoggerConfig as ActSaveLoggerConfig
 from nshtrainer.loggers import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.loggers import CSVLoggerConfig as CSVLoggerConfig
+from nshtrainer.loggers import LoggerConfig as LoggerConfig
 from nshtrainer.loggers import TensorboardLoggerConfig as TensorboardLoggerConfig
 from nshtrainer.loggers import WandbLoggerConfig as WandbLoggerConfig
 from nshtrainer.loggers.wandb import CallbackConfigBase as CallbackConfigBase
@@ -26,6 +27,7 @@ __all__ = [
     "BaseLoggerConfig",
     "CSVLoggerConfig",
     "CallbackConfigBase",
+    "LoggerConfig",
     "TensorboardLoggerConfig",
     "WandbLoggerConfig",
     "WandbUploadCodeCallbackConfig",

nshtrainer/configs/lr_scheduler/__init__.py

@@ -5,8 +5,12 @@ __codegen__ = True
 from nshtrainer.lr_scheduler import (
     LinearWarmupCosineDecayLRSchedulerConfig as LinearWarmupCosineDecayLRSchedulerConfig,
 )
+from nshtrainer.lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
 from nshtrainer.lr_scheduler import LRSchedulerConfigBase as LRSchedulerConfigBase
 from nshtrainer.lr_scheduler import ReduceLROnPlateauConfig as ReduceLROnPlateauConfig
+from nshtrainer.lr_scheduler.linear_warmup_cosine import (
+    DurationConfig as DurationConfig,
+)
 from nshtrainer.lr_scheduler.reduce_lr_on_plateau import MetricConfig as MetricConfig
 
 from . import _base as _base
@@ -14,6 +18,8 @@ from . import linear_warmup_cosine as linear_warmup_cosine
 from . import reduce_lr_on_plateau as reduce_lr_on_plateau
 
 __all__ = [
+    "DurationConfig",
+    "LRSchedulerConfig",
     "LRSchedulerConfigBase",
     "LinearWarmupCosineDecayLRSchedulerConfig",
     "MetricConfig",

nshtrainer/configs/lr_scheduler/linear_warmup_cosine/__init__.py

@@ -2,6 +2,9 @@ from __future__ import annotations
 
 __codegen__ = True
 
+from nshtrainer.lr_scheduler.linear_warmup_cosine import (
+    DurationConfig as DurationConfig,
+)
 from nshtrainer.lr_scheduler.linear_warmup_cosine import (
     LinearWarmupCosineDecayLRSchedulerConfig as LinearWarmupCosineDecayLRSchedulerConfig,
 )
@@ -10,6 +13,7 @@ from nshtrainer.lr_scheduler.linear_warmup_cosine import (
 )
 
 __all__ = [
+    "DurationConfig",
     "LRSchedulerConfigBase",
     "LinearWarmupCosineDecayLRSchedulerConfig",
 ]

nshtrainer/configs/nn/__init__.py

@@ -8,6 +8,7 @@ from nshtrainer.nn import GELUNonlinearityConfig as GELUNonlinearityConfig
 from nshtrainer.nn import LeakyReLUNonlinearityConfig as LeakyReLUNonlinearityConfig
 from nshtrainer.nn import MishNonlinearityConfig as MishNonlinearityConfig
 from nshtrainer.nn import MLPConfig as MLPConfig
+from nshtrainer.nn import NonlinearityConfig as NonlinearityConfig
 from nshtrainer.nn import PReLUConfig as PReLUConfig
 from nshtrainer.nn import ReLUNonlinearityConfig as ReLUNonlinearityConfig
 from nshtrainer.nn import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
@@ -31,6 +32,7 @@ __all__ = [
     "LeakyReLUNonlinearityConfig",
     "MLPConfig",
     "MishNonlinearityConfig",
+    "NonlinearityConfig",
     "PReLUConfig",
     "ReLUNonlinearityConfig",
     "SiLUNonlinearityConfig",

nshtrainer/configs/nn/mlp/__init__.py

@@ -4,8 +4,10 @@ __codegen__ = True
 
 from nshtrainer.nn.mlp import BaseNonlinearityConfig as BaseNonlinearityConfig
 from nshtrainer.nn.mlp import MLPConfig as MLPConfig
+from nshtrainer.nn.mlp import NonlinearityConfig as NonlinearityConfig
 
 __all__ = [
     "BaseNonlinearityConfig",
     "MLPConfig",
+    "NonlinearityConfig",
 ]

nshtrainer/configs/nn/nonlinearity/__init__.py

@@ -9,6 +9,7 @@ from nshtrainer.nn.nonlinearity import (
     LeakyReLUNonlinearityConfig as LeakyReLUNonlinearityConfig,
 )
 from nshtrainer.nn.nonlinearity import MishNonlinearityConfig as MishNonlinearityConfig
+from nshtrainer.nn.nonlinearity import NonlinearityConfig as NonlinearityConfig
 from nshtrainer.nn.nonlinearity import PReLUConfig as PReLUConfig
 from nshtrainer.nn.nonlinearity import ReLUNonlinearityConfig as ReLUNonlinearityConfig
 from nshtrainer.nn.nonlinearity import (
@@ -38,6 +39,7 @@ __all__ = [
     "GELUNonlinearityConfig",
     "LeakyReLUNonlinearityConfig",
     "MishNonlinearityConfig",
+    "NonlinearityConfig",
     "PReLUConfig",
     "ReLUNonlinearityConfig",
     "SiLUNonlinearityConfig",

nshtrainer/configs/optimizer/__init__.py

@@ -3,9 +3,11 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.optimizer import AdamWConfig as AdamWConfig
+from nshtrainer.optimizer import OptimizerConfig as OptimizerConfig
 from nshtrainer.optimizer import OptimizerConfigBase as OptimizerConfigBase
 
 __all__ = [
     "AdamWConfig",
+    "OptimizerConfig",
     "OptimizerConfigBase",
 ]

nshtrainer/configs/profiler/__init__.py

@@ -4,6 +4,7 @@ __codegen__ = True
 
 from nshtrainer.profiler import AdvancedProfilerConfig as AdvancedProfilerConfig
 from nshtrainer.profiler import BaseProfilerConfig as BaseProfilerConfig
+from nshtrainer.profiler import ProfilerConfig as ProfilerConfig
 from nshtrainer.profiler import PyTorchProfilerConfig as PyTorchProfilerConfig
 from nshtrainer.profiler import SimpleProfilerConfig as SimpleProfilerConfig
 
@@ -15,6 +16,7 @@ from . import simple as simple
 __all__ = [
     "AdvancedProfilerConfig",
     "BaseProfilerConfig",
+    "ProfilerConfig",
     "PyTorchProfilerConfig",
     "SimpleProfilerConfig",
     "_base",

nshtrainer/configs/trainer/__init__.py

@@ -9,7 +9,11 @@ from nshtrainer.trainer._config import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.trainer._config import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
+from nshtrainer.trainer._config import CallbackConfig as CallbackConfig
 from nshtrainer.trainer._config import CallbackConfigBase as CallbackConfigBase
+from nshtrainer.trainer._config import (
+    CheckpointCallbackConfig as CheckpointCallbackConfig,
+)
 from nshtrainer.trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
 from nshtrainer.trainer._config import CSVLoggerConfig as CSVLoggerConfig
 from nshtrainer.trainer._config import (
@@ -29,6 +33,7 @@ from nshtrainer.trainer._config import (
     LearningRateMonitorConfig as LearningRateMonitorConfig,
 )
 from nshtrainer.trainer._config import LogEpochCallbackConfig as LogEpochCallbackConfig
+from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
 from nshtrainer.trainer._config import MetricConfig as MetricConfig
 from nshtrainer.trainer._config import (
     NormLoggingCallbackConfig as NormLoggingCallbackConfig,
@@ -37,6 +42,7 @@ from nshtrainer.trainer._config import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
 from nshtrainer.trainer._config import PluginConfigBase as PluginConfigBase
+from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
@@ -48,6 +54,9 @@ from nshtrainer.trainer._config import StrategyConfigBase as StrategyConfigBase
 from nshtrainer.trainer._config import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
+from nshtrainer.trainer._config import (
+    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
+)
 from nshtrainer.trainer._config import WandbLoggerConfig as WandbLoggerConfig
 
 from . import _config as _config
@@ -59,7 +68,9 @@ __all__ = [
     "BaseLoggerConfig",
     "BestCheckpointCallbackConfig",
     "CSVLoggerConfig",
+    "CallbackConfig",
     "CallbackConfigBase",
+    "CheckpointCallbackConfig",
     "CheckpointSavingConfig",
     "DebugFlagCallbackConfig",
     "DirectoryConfig",
@@ -70,15 +81,18 @@ __all__ = [
     "LastCheckpointCallbackConfig",
     "LearningRateMonitorConfig",
     "LogEpochCallbackConfig",
+    "LoggerConfig",
     "MetricConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
     "PluginConfigBase",
+    "ProfilerConfig",
     "RLPSanityChecksCallbackConfig",
     "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
     "StrategyConfigBase",
     "TensorboardLoggerConfig",
+    "TimeCheckpointCallbackConfig",
     "TrainerConfig",
     "WandbLoggerConfig",
     "_config",

nshtrainer/configs/trainer/_config/__init__.py

@@ -8,7 +8,11 @@ from nshtrainer.trainer._config import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.trainer._config import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
+from nshtrainer.trainer._config import CallbackConfig as CallbackConfig
 from nshtrainer.trainer._config import CallbackConfigBase as CallbackConfigBase
+from nshtrainer.trainer._config import (
+    CheckpointCallbackConfig as CheckpointCallbackConfig,
+)
 from nshtrainer.trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
 from nshtrainer.trainer._config import CSVLoggerConfig as CSVLoggerConfig
 from nshtrainer.trainer._config import (
@@ -28,6 +32,7 @@ from nshtrainer.trainer._config import (
     LearningRateMonitorConfig as LearningRateMonitorConfig,
 )
 from nshtrainer.trainer._config import LogEpochCallbackConfig as LogEpochCallbackConfig
+from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
 from nshtrainer.trainer._config import MetricConfig as MetricConfig
 from nshtrainer.trainer._config import (
     NormLoggingCallbackConfig as NormLoggingCallbackConfig,
@@ -36,6 +41,7 @@ from nshtrainer.trainer._config import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
 from nshtrainer.trainer._config import PluginConfigBase as PluginConfigBase
+from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
@@ -47,6 +53,9 @@ from nshtrainer.trainer._config import StrategyConfigBase as StrategyConfigBase
 from nshtrainer.trainer._config import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
+from nshtrainer.trainer._config import (
+    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
+)
 from nshtrainer.trainer._config import TrainerConfig as TrainerConfig
 from nshtrainer.trainer._config import WandbLoggerConfig as WandbLoggerConfig
 
@@ -56,7 +65,9 @@ __all__ = [
     "BaseLoggerConfig",
     "BestCheckpointCallbackConfig",
     "CSVLoggerConfig",
+    "CallbackConfig",
     "CallbackConfigBase",
+    "CheckpointCallbackConfig",
     "CheckpointSavingConfig",
     "DebugFlagCallbackConfig",
     "DirectoryConfig",
@@ -67,15 +78,18 @@ __all__ = [
     "LastCheckpointCallbackConfig",
     "LearningRateMonitorConfig",
     "LogEpochCallbackConfig",
+    "LoggerConfig",
     "MetricConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
     "PluginConfigBase",
+    "ProfilerConfig",
     "RLPSanityChecksCallbackConfig",
     "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
     "StrategyConfigBase",
     "TensorboardLoggerConfig",
+    "TimeCheckpointCallbackConfig",
     "TrainerConfig",
     "WandbLoggerConfig",
 ]

nshtrainer/configs/util/__init__.py

@@ -32,6 +32,7 @@ from nshtrainer.util._environment_info import (
 )
 from nshtrainer.util._environment_info import GitRepositoryConfig as GitRepositoryConfig
 from nshtrainer.util.config import DTypeConfig as DTypeConfig
+from nshtrainer.util.config import DurationConfig as DurationConfig
 from nshtrainer.util.config import EpochsConfig as EpochsConfig
 from nshtrainer.util.config import StepsConfig as StepsConfig
 
@@ -40,6 +41,7 @@ from . import config as config
 
 __all__ = [
     "DTypeConfig",
+    "DurationConfig",
     "EnvironmentCUDAConfig",
     "EnvironmentClassInformationConfig",
     "EnvironmentConfig",

nshtrainer/configs/util/config/__init__.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.util.config import DTypeConfig as DTypeConfig
+from nshtrainer.util.config import DurationConfig as DurationConfig
 from nshtrainer.util.config import EpochsConfig as EpochsConfig
 from nshtrainer.util.config import StepsConfig as StepsConfig
 
@@ -11,6 +12,7 @@ from . import duration as duration
 
 __all__ = [
     "DTypeConfig",
+    "DurationConfig",
     "EpochsConfig",
     "StepsConfig",
     "dtype",

nshtrainer/configs/util/config/duration/__init__.py

@@ -2,10 +2,12 @@ from __future__ import annotations
 
 __codegen__ = True
 
+from nshtrainer.util.config.duration import DurationConfig as DurationConfig
 from nshtrainer.util.config.duration import EpochsConfig as EpochsConfig
 from nshtrainer.util.config.duration import StepsConfig as StepsConfig
 
 __all__ = [
+    "DurationConfig",
     "EpochsConfig",
     "StepsConfig",
 ]

nshtrainer/loggers/__init__.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
-from typing import Annotated, TypeAlias
+from typing import Annotated
 
 import nshconfig as C
+from typing_extensions import TypeAliasType
 
 from ._base import BaseLoggerConfig as BaseLoggerConfig
 from .actsave import ActSaveLoggerConfig as ActSaveLoggerConfig
@@ -10,7 +11,13 @@ from .csv import CSVLoggerConfig as CSVLoggerConfig
 from .tensorboard import TensorboardLoggerConfig as TensorboardLoggerConfig
 from .wandb import WandbLoggerConfig as WandbLoggerConfig
 
-LoggerConfig: TypeAlias = Annotated[
-    CSVLoggerConfig | TensorboardLoggerConfig | WandbLoggerConfig | ActSaveLoggerConfig,
-    C.Field(discriminator="name"),
-]
+LoggerConfig = TypeAliasType(
+    "LoggerConfig",
+    Annotated[
+        CSVLoggerConfig
+        | TensorboardLoggerConfig
+        | WandbLoggerConfig
+        | ActSaveLoggerConfig,
+        C.Field(discriminator="name"),
+    ],
+)

nshtrainer/lr_scheduler/__init__.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
-from typing import Annotated, TypeAlias
+from typing import Annotated
 
 import nshconfig as C
+from typing_extensions import TypeAliasType
 
 from ._base import LRSchedulerConfigBase as LRSchedulerConfigBase
 from ._base import LRSchedulerMetadata as LRSchedulerMetadata
@@ -15,7 +16,10 @@ from .linear_warmup_cosine import (
 from .reduce_lr_on_plateau import ReduceLROnPlateau as ReduceLROnPlateau
 from .reduce_lr_on_plateau import ReduceLROnPlateauConfig as ReduceLROnPlateauConfig
 
-LRSchedulerConfig: TypeAlias = Annotated[
-    LinearWarmupCosineDecayLRSchedulerConfig | ReduceLROnPlateauConfig,
-    C.Field(discriminator="name"),
-]
+LRSchedulerConfig = TypeAliasType(
+    "LRSchedulerConfig",
+    Annotated[
+        LinearWarmupCosineDecayLRSchedulerConfig | ReduceLROnPlateauConfig,
+        C.Field(discriminator="name"),
+    ],
+)

nshtrainer/model/mixins/callback.py

@@ -2,18 +2,20 @@ from __future__ import annotations
 
 import logging
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, TypeAlias, cast
+from typing import Any, cast
 
 from lightning.pytorch import Callback, LightningModule
-from typing_extensions import override
+from typing_extensions import TypeAliasType, override
 
 from ..._callback import NTCallbackBase
 from ...util.typing_utils import mixin_base_type
 
 log = logging.getLogger(__name__)
 
-_Callback = Callback | NTCallbackBase
-CallbackFn: TypeAlias = Callable[[], _Callback | Iterable[_Callback] | None]
+_Callback = TypeAliasType("_Callback", Callback | NTCallbackBase)
+CallbackFn = TypeAliasType(
+    "CallbackFn", Callable[[], _Callback | Iterable[_Callback] | None]
+)
 
 
 class CallbackRegistrarModuleMixin:

nshtrainer/optimizer.py

@@ -2,12 +2,12 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
-from typing import Annotated, Any, Literal, TypeAlias
+from typing import Annotated, Any, Literal
 
 import nshconfig as C
 import torch.nn as nn
 from torch.optim import Optimizer
-from typing_extensions import override
+from typing_extensions import TypeAliasType, override
 
 
 class OptimizerConfigBase(C.Config, ABC):
@@ -57,4 +57,6 @@ class AdamWConfig(OptimizerConfigBase):
         )
 
 
-OptimizerConfig: TypeAlias = Annotated[AdamWConfig, C.Field(discriminator="name")]
+OptimizerConfig = TypeAliasType(
+    "OptimizerConfig", Annotated[AdamWConfig, C.Field(discriminator="name")]
+)

nshtrainer/profiler/__init__.py

@@ -1,15 +1,19 @@
 from __future__ import annotations
 
-from typing import Annotated, TypeAlias
+from typing import Annotated
 
 import nshconfig as C
+from typing_extensions import TypeAliasType
 
 from ._base import BaseProfilerConfig as BaseProfilerConfig
 from .advanced import AdvancedProfilerConfig as AdvancedProfilerConfig
 from .pytorch import PyTorchProfilerConfig as PyTorchProfilerConfig
 from .simple import SimpleProfilerConfig as SimpleProfilerConfig
 
-ProfilerConfig: TypeAlias = Annotated[
-    SimpleProfilerConfig | AdvancedProfilerConfig | PyTorchProfilerConfig,
-    C.Field(discriminator="name"),
-]
+ProfilerConfig = TypeAliasType(
+    "ProfilerConfig",
+    Annotated[
+        SimpleProfilerConfig | AdvancedProfilerConfig | PyTorchProfilerConfig,
+        C.Field(discriminator="name"),
+    ],
+)

nshtrainer/trainer/_config.py

@@ -14,7 +14,6 @@ from typing import (
     Any,
     ClassVar,
     Literal,
-    TypeAlias,
 )
 
 import nshconfig as C
@@ -101,47 +100,53 @@ class StrategyConfigBase(C.Config, ABC):
 strategy_registry = C.Registry(StrategyConfigBase, discriminator="name")
 
 
-AcceleratorLiteral: TypeAlias = Literal[
-    "cpu", "gpu", "tpu", "ipu", "hpu", "mps", "auto"
-]
-
-StrategyLiteral: TypeAlias = Literal[
-    "auto",
-    "ddp",
-    "ddp_find_unused_parameters_false",
-    "ddp_find_unused_parameters_true",
-    "ddp_spawn",
-    "ddp_spawn_find_unused_parameters_false",
-    "ddp_spawn_find_unused_parameters_true",
-    "ddp_fork",
-    "ddp_fork_find_unused_parameters_false",
-    "ddp_fork_find_unused_parameters_true",
-    "ddp_notebook",
-    "dp",
-    "deepspeed",
-    "deepspeed_stage_1",
-    "deepspeed_stage_1_offload",
-    "deepspeed_stage_2",
-    "deepspeed_stage_2_offload",
-    "deepspeed_stage_3",
-    "deepspeed_stage_3_offload",
-    "deepspeed_stage_3_offload_nvme",
-    "fsdp",
-    "fsdp_cpu_offload",
-    "single_xla",
-    "xla_fsdp",
-    "xla",
-    "single_tpu",
-]
-
-
-CheckpointCallbackConfig: TypeAlias = Annotated[
-    BestCheckpointCallbackConfig
-    | LastCheckpointCallbackConfig
-    | OnExceptionCheckpointCallbackConfig
-    | TimeCheckpointCallbackConfig,
-    C.Field(discriminator="name"),
-]
+AcceleratorLiteral = TypeAliasType(
+    "AcceleratorLiteral", Literal["cpu", "gpu", "tpu", "ipu", "hpu", "mps", "auto"]
+)
+
+StrategyLiteral = TypeAliasType(
+    "StrategyLiteral",
+    Literal[
+        "auto",
+        "ddp",
+        "ddp_find_unused_parameters_false",
+        "ddp_find_unused_parameters_true",
+        "ddp_spawn",
+        "ddp_spawn_find_unused_parameters_false",
+        "ddp_spawn_find_unused_parameters_true",
+        "ddp_fork",
+        "ddp_fork_find_unused_parameters_false",
+        "ddp_fork_find_unused_parameters_true",
+        "ddp_notebook",
+        "dp",
+        "deepspeed",
+        "deepspeed_stage_1",
+        "deepspeed_stage_1_offload",
+        "deepspeed_stage_2",
+        "deepspeed_stage_2_offload",
+        "deepspeed_stage_3",
+        "deepspeed_stage_3_offload",
+        "deepspeed_stage_3_offload_nvme",
+        "fsdp",
+        "fsdp_cpu_offload",
+        "single_xla",
+        "xla_fsdp",
+        "xla",
+        "single_tpu",
+    ],
+)
+
+
+CheckpointCallbackConfig = TypeAliasType(
+    "CheckpointCallbackConfig",
+    Annotated[
+        BestCheckpointCallbackConfig
+        | LastCheckpointCallbackConfig
+        | OnExceptionCheckpointCallbackConfig
+        | TimeCheckpointCallbackConfig,
+        C.Field(discriminator="name"),
+    ],
+)
 
 
 class CheckpointSavingConfig(CallbackConfigBase):

nshtrainer/trainer/_runtime_callback.py

@@ -4,14 +4,14 @@ import datetime
 import logging
 import time
 from dataclasses import dataclass
-from typing import Any, Literal, TypeAlias
+from typing import Any, Literal
 
 from lightning.pytorch.callbacks.callback import Callback
-from typing_extensions import override
+from typing_extensions import TypeAliasType, override
 
 log = logging.getLogger(__name__)
 
-Stage: TypeAlias = Literal["train", "validate", "test", "predict"]
+Stage = TypeAliasType("Stage", Literal["train", "validate", "test", "predict"])
 ALL_STAGES = ("train", "validate", "test", "predict")
 
 

nshtrainer/trainer/signal_connector.py

@@ -12,7 +12,7 @@ from collections import defaultdict
 from collections.abc import Callable
 from pathlib import Path
 from types import FrameType
-from typing import Any, TypeAlias
+from typing import Any
 
 import nshrunner as nr
 import torch.utils.data
@@ -22,12 +22,14 @@ from lightning.pytorch.trainer.connectors.signal_connector import _HandlersCompo
 from lightning.pytorch.trainer.connectors.signal_connector import (
     _SignalConnector as _LightningSignalConnector,
 )
-from typing_extensions import override
+from typing_extensions import TypeAliasType, override
 
 log = logging.getLogger(__name__)
 
-_SIGNUM = int | signal.Signals
-_HANDLER: TypeAlias = Callable[[_SIGNUM, FrameType], Any] | int | signal.Handlers | None
+_SIGNUM = TypeAliasType("_SIGNUM", int | signal.Signals)
+_HANDLER = TypeAliasType(
+    "_HANDLER", Callable[[_SIGNUM, FrameType], Any] | int | signal.Handlers | None
+)
 _IS_WINDOWS = platform.system() == "Windows"
 
 

nshtrainer/util/_useful_types.py

@@ -7,7 +7,14 @@ from collections.abc import Set as AbstractSet
 from os import PathLike
 from typing import Any, TypeVar, overload
 
-from typing_extensions import Buffer, Literal, Protocol, SupportsIndex, TypeAlias
+from typing_extensions import (
+    Buffer,
+    Literal,
+    Protocol,
+    SupportsIndex,
+    TypeAlias,
+    TypeAliasType,
+)
 
 _KT = TypeVar("_KT")
 _KT_co = TypeVar("_KT_co", covariant=True)
@@ -60,7 +67,9 @@ class SupportsAllComparisons(
 ): ...
 
 
-SupportsRichComparison: TypeAlias = SupportsDunderLT[Any] | SupportsDunderGT[Any]
+SupportsRichComparison = TypeAliasType(
+    "SupportsRichComparison", SupportsDunderLT[Any] | SupportsDunderGT[Any]
+)
 SupportsRichComparisonT = TypeVar(
     "SupportsRichComparisonT", bound=SupportsRichComparison
 )

nshtrainer/util/config/dtype.py

@@ -1,57 +1,60 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Literal, TypeAlias
+from typing import TYPE_CHECKING, Literal
 
 import nshconfig as C
 import torch
-from typing_extensions import assert_never
+from typing_extensions import TypeAliasType, assert_never
 
 from ..bf16 import is_bf16_supported_no_emulation
 
 if TYPE_CHECKING:
     from ...trainer._config import TrainerConfig
 
-DTypeName: TypeAlias = Literal[
-    "float32",
-    "float",
-    "float64",
-    "double",
-    "float16",
-    "bfloat16",
-    "float8_e4m3fn",
-    "float8_e4m3fnuz",
-    "float8_e5m2",
-    "float8_e5m2fnuz",
-    "half",
-    "uint8",
-    "uint16",
-    "uint32",
-    "uint64",
-    "int8",
-    "int16",
-    "short",
-    "int32",
-    "int",
-    "int64",
-    "long",
-    "complex32",
-    "complex64",
-    "chalf",
-    "cfloat",
-    "complex128",
-    "cdouble",
-    "quint8",
-    "qint8",
-    "qint32",
-    "bool",
-    "quint4x2",
-    "quint2x4",
-    "bits1x8",
-    "bits2x4",
-    "bits4x2",
-    "bits8",
-    "bits16",
-]
+DTypeName = TypeAliasType(
+    "DTypeName",
+    Literal[
+        "float32",
+        "float",
+        "float64",
+        "double",
+        "float16",
+        "bfloat16",
+        "float8_e4m3fn",
+        "float8_e4m3fnuz",
+        "float8_e5m2",
+        "float8_e5m2fnuz",
+        "half",
+        "uint8",
+        "uint16",
+        "uint32",
+        "uint64",
+        "int8",
+        "int16",
+        "short",
+        "int32",
+        "int",
+        "int64",
+        "long",
+        "complex32",
+        "complex64",
+        "chalf",
+        "cfloat",
+        "complex128",
+        "cdouble",
+        "quint8",
+        "qint8",
+        "qint32",
+        "bool",
+        "quint4x2",
+        "quint2x4",
+        "bits1x8",
+        "bits2x4",
+        "bits4x2",
+        "bits8",
+        "bits16",
+    ],
+)
 
 
 class DTypeConfig(C.Config):

nshtrainer/util/path.py

@@ -6,11 +6,12 @@ import os
 import platform
 import shutil
 from pathlib import Path
-from typing import TypeAlias
+
+from typing_extensions import TypeAliasType
 
 log = logging.getLogger(__name__)
 
-_Path: TypeAlias = str | Path | os.PathLike
+_Path = TypeAliasType("_Path", str | Path)
 
 
 def get_relative_path(source: _Path, destination: _Path):

{nshtrainer-1.0.0b26.dist-info → nshtrainer-1.0.0b28.dist-info}/METADATA

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

{nshtrainer-1.0.0b26.dist-info → nshtrainer-1.0.0b28.dist-info}/RECORD

@@ -7,8 +7,8 @@ nshtrainer/_directory.py,sha256=p2uk1FnISFEpMqlDevKhoWhQsCEtvHUPg459K-86QA8,3053
 nshtrainer/_experimental/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 nshtrainer/_hf_hub.py,sha256=xj91CEKcuvpaiipZbVZovZiU_fdsTkPXOkQ-3Xb-FhU,14183
 nshtrainer/callbacks/__init__.py,sha256=z2uDaUVzNXG79UOAYGpwC0hxk__LTbEBVhXJ3W-Srws,3911
-nshtrainer/callbacks/actsave.py,sha256=framZaRes5G-QjWgeDAJ72k3T3tFkj6Z7qmnlIUzLBQ,3759
-nshtrainer/callbacks/base.py,sha256=JZ-WyIKeUYe8g2pLzs648WkxNfkTke7237n1LU_Ujto,3571
+nshtrainer/callbacks/actsave.py,sha256=bnqS3y9sit1hfUvqPx-WWRug2OeTHPM0PDEJaex7f3Q,3776
+nshtrainer/callbacks/base.py,sha256=AFqvKNzH710xviYQ7X0hw0M7ETWRhBAWv3PN9WnrZw0,3608
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=09cwgAJawTFLdzZwYhg_jVgbW-1d09hwjHdI-PQRck0,797
 nshtrainer/callbacks/checkpoint/_base.py,sha256=ZVEUVl5kjCSSe69Q0rMUbKBNNUog0pxBwWkeyuxG2w0,6304
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=hcvtpNvuyWxxP36XiGyFYwVsSwmHphKD-1IA95DinME,2602
@@ -27,23 +27,24 @@ nshtrainer/callbacks/lr_monitor.py,sha256=IyFZoXaxJoTBSkdLu1iEZ1qI8_UFNJwafR_xTV
 nshtrainer/callbacks/norm_logging.py,sha256=C44Mvt73gqQEpCFd0j3qYg6NY7sL2jm3X1qJVY_XLfI,6329
 nshtrainer/callbacks/print_table.py,sha256=lS49Hz0OLcv3VPxEfLBguwe57y2nmKg0pMF6HJxuJio,2974
 nshtrainer/callbacks/rlp_sanity_checks.py,sha256=kWl2dYOXn2L8k6ub_012jNkqOxtyea1yr1qWRNG6UW4,9990
-nshtrainer/callbacks/shared_parameters.py,sha256=33eRzifNj6reKbvmGuam1hUofo3sD4J6Zz8gtoHK93I,2977
+nshtrainer/callbacks/shared_parameters.py,sha256=ggMI1krkqN7sGOrjK_I96IsTMYMXHoVQm7W90LZb9so,3015
 nshtrainer/callbacks/timer.py,sha256=BB-M7tV4QNYOwY_Su6j9P7IILxVRae_upmDq4qsxiao,4670
 nshtrainer/callbacks/wandb_upload_code.py,sha256=PTqNE1QB5U8NR5zhbiQZrmQuugX2UV7B12UdMpo9aV0,2353
 nshtrainer/callbacks/wandb_watch.py,sha256=tTTcFzxd2Ia9xu8tCogQ5CLJZBq1ne5JlpGVE75vKYs,2976
-nshtrainer/configs/__init__.py,sha256=83LkGk9FuaIpIkxrJKylELAokX4p8u9ZmdeAdG__06E,9681
+nshtrainer/configs/__init__.py,sha256=zyo4lV9ObB3T3_hhBhzWGNb6MRma4h7QHD3OrypxqEw,10582
 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=7H3fIh9c31ce0r8JpuzEY8bZptI7tiVLNwVtj729HAY,303
+nshtrainer/configs/_directory/__init__.py,sha256=_oO7vM9DhzHSxtZcv86sTi7hZIptnK1gr-AP9mqQ370,386
 nshtrainer/configs/_hf_hub/__init__.py,sha256=VUgQnyEI2ekBxBIV15L09tKdrfGt7eWxnf30DiCLaso,416
-nshtrainer/configs/callbacks/__init__.py,sha256=muhbnqVN-KkcdmW5dRJwNMoDGwxURIMt9ggVO1qjea8,4126
+nshtrainer/configs/callbacks/__init__.py,sha256=APpF2jmafqbS4CoMmDvFADi0wdmXJ_BvFw4QnnQpok0,4353
 nshtrainer/configs/callbacks/actsave/__init__.py,sha256=JJg9d8iNGpO-9M1LsK4h1cu3NYWniyIyLQ4SauFCzOs,272
 nshtrainer/configs/callbacks/base/__init__.py,sha256=V694hzF_ubnA-hwTps30PeFbgDSm3I_UIMTnljM3_OI,176
-nshtrainer/configs/callbacks/checkpoint/__init__.py,sha256=OOMwYG4zjvs54hrTyGe2a5MdOjCpTbtpk1wStraZLcA,1342
+nshtrainer/configs/callbacks/checkpoint/__init__.py,sha256=zPUItxoYWrMT9i1TxOvVhIeTa0NEFg-nDE5FjHfkP-A,1564
 nshtrainer/configs/callbacks/checkpoint/_base/__init__.py,sha256=5jl6A5Gv6arZXmHV6lz5dQ8DL6PdJIfJqHLP4acClKQ,479
 nshtrainer/configs/callbacks/checkpoint/best_checkpoint/__init__.py,sha256=_KTkF3_Yx0WiwOvRf2s1KRvod2dryeGJITVkx10YqBE,648
 nshtrainer/configs/callbacks/checkpoint/last_checkpoint/__init__.py,sha256=sH4aJyovCeT6h4xz9r5WfVA0eviJur4zZt7R0hQsyAk,539
 nshtrainer/configs/callbacks/checkpoint/on_exception_checkpoint/__init__.py,sha256=iUoCrTJpvDGEYCYfNCpkficH6D7yq119hyxy1qeFSGU,410
+nshtrainer/configs/callbacks/checkpoint/time_checkpoint/__init__.py,sha256=ZQz5gXKSHz7xvnZbgU_bauNhAJDh9rblBA0ldkSsdwg,539
 nshtrainer/configs/callbacks/debug_flag/__init__.py,sha256=gPC3EAqzuyP2hAcCf3s09sPDe7q_02S1eUCWvRTNKrI,317
 nshtrainer/configs/callbacks/directory_setup/__init__.py,sha256=25551zMMctAkzcLEGBN7HSeQUIrBtbBq7whgPZjtepY,351
 nshtrainer/configs/callbacks/early_stopping/__init__.py,sha256=Q-hAAIcucLNriw5PXIshgBv7Yr5sMDk4GwjTIDdWBxo,434
@@ -59,46 +60,46 @@ nshtrainer/configs/callbacks/shared_parameters/__init__.py,sha256=Ivef5jk3RMYQDe
 nshtrainer/configs/callbacks/timer/__init__.py,sha256=RHOQoREp4NxS_AvKNdc0UuUlS0PnqCxxsuOz5D8h7iM,310
 nshtrainer/configs/callbacks/wandb_upload_code/__init__.py,sha256=WM9hCGFl2LXDUOgkIGaV3tkdnXnVBasrhIILjbIeFUo,358
 nshtrainer/configs/callbacks/wandb_watch/__init__.py,sha256=MW-ANrF529DxBhopovPjYEQ7nANX9ttd1K4_bJnKXks,322
-nshtrainer/configs/loggers/__init__.py,sha256=WcP-g5dgP_cgLHtv7G2gG72nwq_7KzExSiz9ioQ1SJw,1171
+nshtrainer/configs/loggers/__init__.py,sha256=5wTekL79mQxit8f1K3AMllvb0mKertTzOKfC3gpE2Zk,1251
 nshtrainer/configs/loggers/_base/__init__.py,sha256=HxPPPePsEjlNuhnjsMgYIl0rwj_iqNKKOBTEk_zIOsM,169
 nshtrainer/configs/loggers/actsave/__init__.py,sha256=2lZQ4bpbjwd4MuUE_Z_PGbmQjjGtWCZUCtXqKO4dTSc,280
 nshtrainer/configs/loggers/csv/__init__.py,sha256=M3QGF5GKiRGENy3re6LJKpa4A4RThy1FlmaFuR4cPyo,260
 nshtrainer/configs/loggers/tensorboard/__init__.py,sha256=FbkYXnSohIX6JN5XyI-9y91IJv_T3VB3IwmpagXAnM4,309
 nshtrainer/configs/loggers/wandb/__init__.py,sha256=76qb0HhWojf0Ub1x9OkMjtzeXxE67KysBGa-MBbJyC4,651
-nshtrainer/configs/lr_scheduler/__init__.py,sha256=19rkQMI8j31SVp-LAmmV3w9J2Lpv0XdyEzoeQT2dPfE,802
+nshtrainer/configs/lr_scheduler/__init__.py,sha256=8ORO-QC12SjZ2F_reMoDgr8-O8nxZxX0IKU4fl-cC3A,1023
 nshtrainer/configs/lr_scheduler/_base/__init__.py,sha256=fvGjkUJ1K2RVXjXror22QOtEa-xWFJz2Cz3HrBC5XfA,189
-nshtrainer/configs/lr_scheduler/linear_warmup_cosine/__init__.py,sha256=8-KVX8cBJspiELZAbEdJpp8zTkXVik6mn-LNo_Qv27I,412
+nshtrainer/configs/lr_scheduler/linear_warmup_cosine/__init__.py,sha256=i8LeZh0c4wqtZ1ehZb2LCq7kwOL0OyswMMOnwyI6R04,533
 nshtrainer/configs/lr_scheduler/reduce_lr_on_plateau/__init__.py,sha256=lpXEFZY4cM3znZqYG9IZ1xNNtzttt8VVspSuOz0fb-k,467
 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=2nSgm07ym2D4yiMQ15pyyb26laM6WWt5KM44PnCDK5A,1864
-nshtrainer/configs/nn/mlp/__init__.py,sha256=ZHQX44z6A_DCGYGYd0N_AzaHguiVBMwexMBJD6TjppQ,250
-nshtrainer/configs/nn/nonlinearity/__init__.py,sha256=GnJmCX0CwH6GLbs1RgXYrtmxfcmbD0GlujrxYIRr4ms,1934
-nshtrainer/configs/optimizer/__init__.py,sha256=2N9LK2Dm4w45Ff0tlSG4Z52Ne4OAujNT5s6_K-bE-qA,253
-nshtrainer/configs/profiler/__init__.py,sha256=AGkoWizikpXOKE3YQ5wSBrHLHAV3ZNf5rSl8YuEb84s,681
+nshtrainer/configs/nn/__init__.py,sha256=3hVc81Gs9AJYVkrwJkQ_ye7tLU2HOLdBj-mMkXx2c_I,1957
+nshtrainer/configs/nn/mlp/__init__.py,sha256=eMECrgz-My9mFS7lpWVI3dj1ApB-E7xwfmNc37hUsPI,347
+nshtrainer/configs/nn/nonlinearity/__init__.py,sha256=Gjr2HCx8jJTcfu7sLgn54o2ucGKaBea4encm4AWpKNY,2040
+nshtrainer/configs/optimizer/__init__.py,sha256=IMEsEbiVFXSkj6WmDjNjmKQuRspphs5xZnYZ2gYE39Y,344
+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
 nshtrainer/configs/profiler/pytorch/__init__.py,sha256=soAU1s2_Pa1na4gW8CK-iysJBO5M_7YeZC2_x40iEdg,294
 nshtrainer/configs/profiler/simple/__init__.py,sha256=3Wb11lPuFuyasq8xS1CZ4WLuBCLS_nVSQGVllvOOi0Y,289
-nshtrainer/configs/trainer/__init__.py,sha256=3fFvvgTo6MVR7871d0AXTWd9YrabjpJ_aS3iZP-yu2M,3382
-nshtrainer/configs/trainer/_config/__init__.py,sha256=uYB_e6Dg7GIlPjabh74k34N86DaFrYaM_rIjC12WJ_s,3293
+nshtrainer/configs/trainer/__init__.py,sha256=KIDYjJsc-WYXKiH2RNzAZJD5MKOTdO9wdtu_vWDNPxU,3936
+nshtrainer/configs/trainer/_config/__init__.py,sha256=1_Ad5uTvXdVuHMJB3s8s-0EraDwNZssg3sXBmVouF9w,3847
 nshtrainer/configs/trainer/trainer/__init__.py,sha256=DDuBRx0kVNMW0z_sqKTUt8-Ql7bOpargi4KcHHvDu_c,486
-nshtrainer/configs/util/__init__.py,sha256=gtYtZ4VGwEvF9_hByZl8CWOSeDpEOIkkcLtUwvNbSEQ,2014
+nshtrainer/configs/util/__init__.py,sha256=qXittS7f7MyaqJnjvFLKnKsyb6bXTD3dEV16jXVDaH4,2104
 nshtrainer/configs/util/_environment_info/__init__.py,sha256=eB4E0Ck7XCeSC5gbUdA5thd7TXnjGCL0t8GZIFj7uCI,1644
-nshtrainer/configs/util/config/__init__.py,sha256=GHQZR-M0LwL7Qow2oCgmaWwz9h16NkfiWpxIT9cF52Y,411
+nshtrainer/configs/util/config/__init__.py,sha256=nEFiDG3-dvvTytYn1tEkPFzp7fgaGRp2j7toSN7yRGs,501
 nshtrainer/configs/util/config/dtype/__init__.py,sha256=PmGF-O4r6SXqEaagVsQ5YxEqhdVdcU0dgJW1Ljzpp6k,158
-nshtrainer/configs/util/config/duration/__init__.py,sha256=rja-dB-WC2criHrSBC7gkl5GnWeXQ3bHD48zy9kEPbo,254
+nshtrainer/configs/util/config/duration/__init__.py,sha256=44lS2irOIPVfgshMTfnZM2jC6l0Pjst9w2M_lJoS_MU,353
 nshtrainer/data/__init__.py,sha256=K4i3Tw4g9EOK2zlMMbidi99y0SyI4w8P7_XUf1n42Ts,260
 nshtrainer/data/balanced_batch_sampler.py,sha256=r1cBKRXKHD8E1Ax6tj-FUbE-z1qpbO58mQ9VrK9uLnc,5481
 nshtrainer/data/datamodule.py,sha256=lSOgH32nysJWa6Y7ba1QyOdUV0DVVdO98qokP8wigjk,4138
 nshtrainer/data/transform.py,sha256=qd0lIocO59Fk_m90xyOHgFezbymd1mRwly8nbYIfHGc,2263
-nshtrainer/loggers/__init__.py,sha256=11X6D_lF0cHAkxkYsVZY3Q3r30Fq0FUi9heeb5RD870,570
+nshtrainer/loggers/__init__.py,sha256=-y8B-9TF6vJdZUQewJNDcZ2aOv04FEUFtKwaiDobIO0,670
 nshtrainer/loggers/_base.py,sha256=nw4AZzJP3Z-fljgQlgq7FkuMkPmYKTsXj7OfJJSmtXI,811
 nshtrainer/loggers/actsave.py,sha256=23Kre-mq-Y9Iw1SRyGmHnHK1bc_0gTWFXJViv9bkVz0,1324
 nshtrainer/loggers/csv.py,sha256=Deh5gm3oROJbQzigV4SHni5JRwSrBdm-4YD3yrcGnHo,1104
 nshtrainer/loggers/tensorboard.py,sha256=jP9V4nlq_MXUaoD6xv1Cws2ioft83Lm8yUJhhGhuUrQ,2268
 nshtrainer/loggers/wandb.py,sha256=EjKQQznLSUSCWO7uIviz9g0dVW4ZLxb_8UVhY4vR7r0,6800
-nshtrainer/lr_scheduler/__init__.py,sha256=o7eDaj-KOetQvEwMjGKpcQofs1jimqUzHbJ5YM5cOMY,772
+nshtrainer/lr_scheduler/__init__.py,sha256=BGnO-okUTZOtF15-UmQ05U4oEatSF5VNs3YeidNEWn4,853
 nshtrainer/lr_scheduler/_base.py,sha256=EhA2f_WiZ79RcXL2nJbwCwNK620c8ugEVUmJ8CcVpsE,3723
 nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=gvUuv031lvWdXboDeH7iAF3ZgNPQK40bQwfmqb11TNk,5492
 nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=vXH5S26ESHO_LPPqW8aDC3S5NGoZYkXeFjAOgttaUX8,2870
@@ -106,7 +107,7 @@ nshtrainer/metrics/__init__.py,sha256=Nqkn_jsDf3n5WtfMcnaaEftYjIIT2b-S7rmsB1MOMk
 nshtrainer/metrics/_config.py,sha256=kR9OJLZXIaZpG8UpHG3EQwOL36HB8yPk6afYjoIM0XM,1324
 nshtrainer/model/__init__.py,sha256=3G-bwPPSRStWdsdwG9-rn0bXcRpEiP1BiQpF_qavtls,97
 nshtrainer/model/base.py,sha256=JL3AmH17GQjQIoMrZl3O0vUI7dj5ZsO5iEJgoLPyzHw,10356
-nshtrainer/model/mixins/callback.py,sha256=Ea_legORzs0N078j0N9RJivDVeWH5KtXDpdJS75IwIo,3098
+nshtrainer/model/mixins/callback.py,sha256=0LPgve4VszHbLipid4mpI1qnnmdGS2spivs0dXLvqHw,3154
 nshtrainer/model/mixins/debug.py,sha256=1LX9KzeFX9JDPs_a6YCdYDZXLhEk_5rBO2aCqlfBy7w,2087
 nshtrainer/model/mixins/logger.py,sha256=27H99FuLaxc6_dDLG2pid4E_5E0-eLGnc2Ifpt0HYIM,6066
 nshtrainer/nn/__init__.py,sha256=sANhrZpeN5syLKOsmXMwhaFl2SBFPWcLaEe1EH22TWQ,1463
@@ -114,29 +115,29 @@ nshtrainer/nn/mlp.py,sha256=2W8bzE96DzCMzGm6WPiPhNFQfhqaoG3GXPn_oKBnlUM,5988
 nshtrainer/nn/module_dict.py,sha256=9plb8aQUx5TUEPhX5jI9u8LrpTeKe7jZAHi8iIqcN8w,2365
 nshtrainer/nn/module_list.py,sha256=UB43pcwD_3nUke_DyLQt-iXKhWdKM6Zjm84lRC1hPYA,1755
 nshtrainer/nn/nonlinearity.py,sha256=mp5XvXRHURB6jwuZ0YyTj5ZoHJYNJNgO2aLtUY1D-2Y,6114
-nshtrainer/optimizer.py,sha256=73UWfuTU4moDQNSKI8VSQSDc6_i7mBtemxvhBkM-fC4,1563
-nshtrainer/profiler/__init__.py,sha256=X6E4cbIRaXHCDgQtT3LFPw5xIRBhqnJ03SVFbsD-Yd8,516
+nshtrainer/optimizer.py,sha256=wmSRpSoU59rstj2RBoifQ15ZwRInYpm0tDBQZ1gqOfE,1596
+nshtrainer/profiler/__init__.py,sha256=RjaNBoVcTFu8lF0dNlFp-2LaPYdonoIbDy2_KhgF0Ek,594
 nshtrainer/profiler/_base.py,sha256=kFcSVn9gJuMwgDxbfyHh46CmEAIPZjxw3yjPbKgzvwA,950
 nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5NOQ,1160
 nshtrainer/profiler/pytorch.py,sha256=8K37XvPnCApUpIK8tA2zNMFIaIiTLSoxKQoiyCPBm1Q,2757
 nshtrainer/profiler/simple.py,sha256=PimjqcU-JuS-8C0ZGHAdwCxgNLij4x0FH6WXsjBQzZs,1005
 nshtrainer/trainer/__init__.py,sha256=MmoydVS6aYeav7zgDAUHxAQrV_PMQsbnZTCuPnLH9Wk,128
-nshtrainer/trainer/_config.py,sha256=ycuPsBP08X-fz8oHcm1UaGsovGuNnVcOHMxU_AH85_s,35036
-nshtrainer/trainer/_runtime_callback.py,sha256=T3epaj1YeIN0R8CS2cg5HNJIB21TyaD_PVNNOPJ6nJs,4200
-nshtrainer/trainer/signal_connector.py,sha256=YMJf6vTnW0JcnBkuYikm9x_9XscaokrCEzCn4THOGao,10776
+nshtrainer/trainer/_config.py,sha256=Mz9J2ZFqxTlttnRA1eScGRgSAuf3-o3i9-xjN7eTm-k,35256
+nshtrainer/trainer/_runtime_callback.py,sha256=6F2Gq27Q8OFfN3RtdNC6QRA8ac0LC1hh4DUE3V5WgbI,4217
+nshtrainer/trainer/signal_connector.py,sha256=GhfGcSzfaTNhnj2QFkBDq5aT7FqbLMA7eC8SYQs8_8w,10828
 nshtrainer/trainer/trainer.py,sha256=HHqT83zWtYY9g5yD6X9aWrVh5VSpILW8PhoE6fp4snE,20734
 nshtrainer/util/_environment_info.py,sha256=MT8mBe6ZolRfKiwU-les1P-lPNPqXpHQcfADrh_A3uY,24629
-nshtrainer/util/_useful_types.py,sha256=dwZokFkIe7M5i2GR3nQ9A1lhGw06DMAFfH5atyquqSA,8000
+nshtrainer/util/_useful_types.py,sha256=7yd1ajSmjwfmZdBPlHVrIG3iXl1-T3n83JI53N8C7as,8080
 nshtrainer/util/bf16.py,sha256=9QhHZCkYSfYpIcxwAMoXyuh2yTSHBzT-EdLQB297jEs,762
 nshtrainer/util/config/__init__.py,sha256=Z39JJufSb61Lhn2GfVcv3eFW_eorOrN9-9llDWlnZZM,272
-nshtrainer/util/config/dtype.py,sha256=VQ02dbouSehCdziFN7NfsEu7o_dcW5SNbl5bhNIwn30,2021
+nshtrainer/util/config/dtype.py,sha256=Fn_MhhQoHPyFAnFPSwvcvLiGR3yWFIszMba02CJiC4g,2213
 nshtrainer/util/config/duration.py,sha256=mM-UfU_HvhXwW33TYEDg0x58n80tnle2e6VaWtxZTjk,764
 nshtrainer/util/environment.py,sha256=s-B5nY0cKYXdFMdNYumvC_xxacMATiI4DvV2gUDu20k,4195
-nshtrainer/util/path.py,sha256=mfuMRk29YpmwdRsyyf99VT0n_T9caM99SS7VfA6oIHE,3501
+nshtrainer/util/path.py,sha256=L-Nh9tlXSUfoP19TFbQq8I0AfS5ugCfGYTYFeddDHcs,3516
 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.0b26.dist-info/METADATA,sha256=9zqciwuzpDhIIKqlwo3IviLhHUDwkDbYJCwbzZX8VQw,988
-nshtrainer-1.0.0b26.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-nshtrainer-1.0.0b26.dist-info/RECORD,,
+nshtrainer-1.0.0b28.dist-info/METADATA,sha256=1MJi65pa7HEVmtDR64Y32SwDe_bv1AZHSgyo6gIBmzo,988
+nshtrainer-1.0.0b28.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+nshtrainer-1.0.0b28.dist-info/RECORD,,