nshtrainer 1.0.0b25__py3-none-any.whl → 1.0.0b27__py3-none-any.whl

nshtrainer/configs/callbacks/directory_setup/__init__.py

@@ -8,3 +8,8 @@ from nshtrainer.callbacks.directory_setup import (
 from nshtrainer.callbacks.directory_setup import (
     DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "DirectorySetupCallbackConfig",
+]

nshtrainer/configs/callbacks/early_stopping/__init__.py

@@ -7,3 +7,9 @@ from nshtrainer.callbacks.early_stopping import (
     EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
 )
 from nshtrainer.callbacks.early_stopping import MetricConfig as MetricConfig
+
+__all__ = [
+    "CallbackConfigBase",
+    "EarlyStoppingCallbackConfig",
+    "MetricConfig",
+]

nshtrainer/configs/callbacks/ema/__init__.py

@@ -4,3 +4,8 @@ __codegen__ = True
 
 from nshtrainer.callbacks.ema import CallbackConfigBase as CallbackConfigBase
 from nshtrainer.callbacks.ema import EMACallbackConfig as EMACallbackConfig
+
+__all__ = [
+    "CallbackConfigBase",
+    "EMACallbackConfig",
+]

nshtrainer/configs/callbacks/finite_checks/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.callbacks.finite_checks import CallbackConfigBase as CallbackCon
 from nshtrainer.callbacks.finite_checks import (
     FiniteChecksCallbackConfig as FiniteChecksCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "FiniteChecksCallbackConfig",
+]

nshtrainer/configs/callbacks/gradient_skipping/__init__.py

@@ -8,3 +8,8 @@ from nshtrainer.callbacks.gradient_skipping import (
 from nshtrainer.callbacks.gradient_skipping import (
     GradientSkippingCallbackConfig as GradientSkippingCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "GradientSkippingCallbackConfig",
+]

nshtrainer/configs/callbacks/log_epoch/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.callbacks.log_epoch import CallbackConfigBase as CallbackConfigB
 from nshtrainer.callbacks.log_epoch import (
     LogEpochCallbackConfig as LogEpochCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "LogEpochCallbackConfig",
+]

nshtrainer/configs/callbacks/lr_monitor/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.callbacks.lr_monitor import CallbackConfigBase as CallbackConfig
 from nshtrainer.callbacks.lr_monitor import (
     LearningRateMonitorConfig as LearningRateMonitorConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "LearningRateMonitorConfig",
+]

nshtrainer/configs/callbacks/norm_logging/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.callbacks.norm_logging import CallbackConfigBase as CallbackConf
 from nshtrainer.callbacks.norm_logging import (
     NormLoggingCallbackConfig as NormLoggingCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "NormLoggingCallbackConfig",
+]

nshtrainer/configs/callbacks/print_table/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.callbacks.print_table import CallbackConfigBase as CallbackConfi
 from nshtrainer.callbacks.print_table import (
     PrintTableMetricsCallbackConfig as PrintTableMetricsCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "PrintTableMetricsCallbackConfig",
+]

nshtrainer/configs/callbacks/rlp_sanity_checks/__init__.py

@@ -8,3 +8,8 @@ from nshtrainer.callbacks.rlp_sanity_checks import (
 from nshtrainer.callbacks.rlp_sanity_checks import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "RLPSanityChecksCallbackConfig",
+]

nshtrainer/configs/callbacks/shared_parameters/__init__.py

@@ -8,3 +8,8 @@ from nshtrainer.callbacks.shared_parameters import (
 from nshtrainer.callbacks.shared_parameters import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "SharedParametersCallbackConfig",
+]

nshtrainer/configs/callbacks/timer/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.callbacks.timer import CallbackConfigBase as CallbackConfigBase
 from nshtrainer.callbacks.timer import (
     EpochTimerCallbackConfig as EpochTimerCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "EpochTimerCallbackConfig",
+]

nshtrainer/configs/callbacks/wandb_upload_code/__init__.py

@@ -8,3 +8,8 @@ from nshtrainer.callbacks.wandb_upload_code import (
 from nshtrainer.callbacks.wandb_upload_code import (
     WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "WandbUploadCodeCallbackConfig",
+]

nshtrainer/configs/callbacks/wandb_watch/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.callbacks.wandb_watch import CallbackConfigBase as CallbackConfi
 from nshtrainer.callbacks.wandb_watch import (
     WandbWatchCallbackConfig as WandbWatchCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "WandbWatchCallbackConfig",
+]

nshtrainer/configs/loggers/__init__.py

@@ -5,7 +5,6 @@ __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
@@ -21,3 +20,19 @@ from . import actsave as actsave
 from . import csv as csv
 from . import tensorboard as tensorboard
 from . import wandb as wandb
+
+__all__ = [
+    "ActSaveLoggerConfig",
+    "BaseLoggerConfig",
+    "CSVLoggerConfig",
+    "CallbackConfigBase",
+    "TensorboardLoggerConfig",
+    "WandbLoggerConfig",
+    "WandbUploadCodeCallbackConfig",
+    "WandbWatchCallbackConfig",
+    "_base",
+    "actsave",
+    "csv",
+    "tensorboard",
+    "wandb",
+]

nshtrainer/configs/loggers/_base/__init__.py

@@ -3,3 +3,7 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.loggers._base import BaseLoggerConfig as BaseLoggerConfig
+
+__all__ = [
+    "BaseLoggerConfig",
+]

nshtrainer/configs/loggers/actsave/__init__.py

@@ -4,3 +4,8 @@ __codegen__ = True
 
 from nshtrainer.loggers.actsave import ActSaveLoggerConfig as ActSaveLoggerConfig
 from nshtrainer.loggers.actsave import BaseLoggerConfig as BaseLoggerConfig
+
+__all__ = [
+    "ActSaveLoggerConfig",
+    "BaseLoggerConfig",
+]

nshtrainer/configs/loggers/csv/__init__.py

@@ -4,3 +4,8 @@ __codegen__ = True
 
 from nshtrainer.loggers.csv import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.loggers.csv import CSVLoggerConfig as CSVLoggerConfig
+
+__all__ = [
+    "BaseLoggerConfig",
+    "CSVLoggerConfig",
+]

nshtrainer/configs/loggers/tensorboard/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.loggers.tensorboard import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.loggers.tensorboard import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
+
+__all__ = [
+    "BaseLoggerConfig",
+    "TensorboardLoggerConfig",
+]

nshtrainer/configs/loggers/wandb/__init__.py

@@ -11,3 +11,11 @@ from nshtrainer.loggers.wandb import (
 from nshtrainer.loggers.wandb import (
     WandbWatchCallbackConfig as WandbWatchCallbackConfig,
 )
+
+__all__ = [
+    "BaseLoggerConfig",
+    "CallbackConfigBase",
+    "WandbLoggerConfig",
+    "WandbUploadCodeCallbackConfig",
+    "WandbWatchCallbackConfig",
+]

nshtrainer/configs/lr_scheduler/__init__.py

@@ -5,14 +5,20 @@ __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
 from . import linear_warmup_cosine as linear_warmup_cosine
 from . import reduce_lr_on_plateau as reduce_lr_on_plateau
+
+__all__ = [
+    "LRSchedulerConfigBase",
+    "LinearWarmupCosineDecayLRSchedulerConfig",
+    "MetricConfig",
+    "ReduceLROnPlateauConfig",
+    "_base",
+    "linear_warmup_cosine",
+    "reduce_lr_on_plateau",
+]

nshtrainer/configs/lr_scheduler/_base/__init__.py

@@ -3,3 +3,7 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.lr_scheduler._base import LRSchedulerConfigBase as LRSchedulerConfigBase
+
+__all__ = [
+    "LRSchedulerConfigBase",
+]

nshtrainer/configs/lr_scheduler/linear_warmup_cosine/__init__.py

@@ -2,12 +2,14 @@ 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,
 )
 from nshtrainer.lr_scheduler.linear_warmup_cosine import (
     LRSchedulerConfigBase as LRSchedulerConfigBase,
 )
+
+__all__ = [
+    "LRSchedulerConfigBase",
+    "LinearWarmupCosineDecayLRSchedulerConfig",
+]

nshtrainer/configs/lr_scheduler/reduce_lr_on_plateau/__init__.py

@@ -9,3 +9,9 @@ from nshtrainer.lr_scheduler.reduce_lr_on_plateau import MetricConfig as MetricC
 from nshtrainer.lr_scheduler.reduce_lr_on_plateau import (
     ReduceLROnPlateauConfig as ReduceLROnPlateauConfig,
 )
+
+__all__ = [
+    "LRSchedulerConfigBase",
+    "MetricConfig",
+    "ReduceLROnPlateauConfig",
+]

nshtrainer/configs/metrics/__init__.py

@@ -5,3 +5,8 @@ __codegen__ = True
 from nshtrainer.metrics import MetricConfig as MetricConfig
 
 from . import _config as _config
+
+__all__ = [
+    "MetricConfig",
+    "_config",
+]

nshtrainer/configs/metrics/_config/__init__.py

@@ -3,3 +3,7 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.metrics._config import MetricConfig as MetricConfig
+
+__all__ = [
+    "MetricConfig",
+]

nshtrainer/configs/nn/__init__.py

@@ -8,7 +8,6 @@ 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
@@ -24,3 +23,24 @@ from nshtrainer.nn.nonlinearity import (
 
 from . import mlp as mlp
 from . import nonlinearity as nonlinearity
+
+__all__ = [
+    "BaseNonlinearityConfig",
+    "ELUNonlinearityConfig",
+    "GELUNonlinearityConfig",
+    "LeakyReLUNonlinearityConfig",
+    "MLPConfig",
+    "MishNonlinearityConfig",
+    "PReLUConfig",
+    "ReLUNonlinearityConfig",
+    "SiLUNonlinearityConfig",
+    "SigmoidNonlinearityConfig",
+    "SoftmaxNonlinearityConfig",
+    "SoftplusNonlinearityConfig",
+    "SoftsignNonlinearityConfig",
+    "SwiGLUNonlinearityConfig",
+    "SwishNonlinearityConfig",
+    "TanhNonlinearityConfig",
+    "mlp",
+    "nonlinearity",
+]

nshtrainer/configs/nn/mlp/__init__.py

@@ -4,4 +4,8 @@ __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",
+]

nshtrainer/configs/nn/nonlinearity/__init__.py

@@ -9,7 +9,6 @@ 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 (
@@ -32,3 +31,21 @@ from nshtrainer.nn.nonlinearity import (
     SwishNonlinearityConfig as SwishNonlinearityConfig,
 )
 from nshtrainer.nn.nonlinearity import TanhNonlinearityConfig as TanhNonlinearityConfig
+
+__all__ = [
+    "BaseNonlinearityConfig",
+    "ELUNonlinearityConfig",
+    "GELUNonlinearityConfig",
+    "LeakyReLUNonlinearityConfig",
+    "MishNonlinearityConfig",
+    "PReLUConfig",
+    "ReLUNonlinearityConfig",
+    "SiLUNonlinearityConfig",
+    "SigmoidNonlinearityConfig",
+    "SoftmaxNonlinearityConfig",
+    "SoftplusNonlinearityConfig",
+    "SoftsignNonlinearityConfig",
+    "SwiGLUNonlinearityConfig",
+    "SwishNonlinearityConfig",
+    "TanhNonlinearityConfig",
+]

nshtrainer/configs/optimizer/__init__.py

@@ -3,5 +3,9 @@ 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",
+    "OptimizerConfigBase",
+]

nshtrainer/configs/profiler/__init__.py

@@ -4,7 +4,6 @@ __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
 
@@ -12,3 +11,14 @@ from . import _base as _base
 from . import advanced as advanced
 from . import pytorch as pytorch
 from . import simple as simple
+
+__all__ = [
+    "AdvancedProfilerConfig",
+    "BaseProfilerConfig",
+    "PyTorchProfilerConfig",
+    "SimpleProfilerConfig",
+    "_base",
+    "advanced",
+    "pytorch",
+    "simple",
+]

nshtrainer/configs/profiler/_base/__init__.py

@@ -3,3 +3,7 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.profiler._base import BaseProfilerConfig as BaseProfilerConfig
+
+__all__ = [
+    "BaseProfilerConfig",
+]

nshtrainer/configs/profiler/advanced/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.profiler.advanced import (
     AdvancedProfilerConfig as AdvancedProfilerConfig,
 )
 from nshtrainer.profiler.advanced import BaseProfilerConfig as BaseProfilerConfig
+
+__all__ = [
+    "AdvancedProfilerConfig",
+    "BaseProfilerConfig",
+]

nshtrainer/configs/profiler/pytorch/__init__.py

@@ -4,3 +4,8 @@ __codegen__ = True
 
 from nshtrainer.profiler.pytorch import BaseProfilerConfig as BaseProfilerConfig
 from nshtrainer.profiler.pytorch import PyTorchProfilerConfig as PyTorchProfilerConfig
+
+__all__ = [
+    "BaseProfilerConfig",
+    "PyTorchProfilerConfig",
+]

nshtrainer/configs/profiler/simple/__init__.py

@@ -4,3 +4,8 @@ __codegen__ = True
 
 from nshtrainer.profiler.simple import BaseProfilerConfig as BaseProfilerConfig
 from nshtrainer.profiler.simple import SimpleProfilerConfig as SimpleProfilerConfig
+
+__all__ = [
+    "BaseProfilerConfig",
+    "SimpleProfilerConfig",
+]

nshtrainer/configs/trainer/__init__.py

@@ -3,16 +3,13 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.trainer import TrainerConfig as TrainerConfig
+from nshtrainer.trainer._config import AcceleratorConfigBase as AcceleratorConfigBase
 from nshtrainer.trainer._config import ActSaveLoggerConfig as ActSaveLoggerConfig
 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 (
@@ -32,7 +29,6 @@ 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,
@@ -40,7 +36,7 @@ from nshtrainer.trainer._config import (
 from nshtrainer.trainer._config import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
-from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
+from nshtrainer.trainer._config import PluginConfigBase as PluginConfigBase
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
@@ -48,10 +44,47 @@ from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingCon
 from nshtrainer.trainer._config import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
+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
 from . import trainer as trainer
+
+__all__ = [
+    "AcceleratorConfigBase",
+    "ActSaveLoggerConfig",
+    "BaseLoggerConfig",
+    "BestCheckpointCallbackConfig",
+    "CSVLoggerConfig",
+    "CallbackConfigBase",
+    "CheckpointSavingConfig",
+    "DebugFlagCallbackConfig",
+    "DirectoryConfig",
+    "EarlyStoppingCallbackConfig",
+    "EnvironmentConfig",
+    "GradientClippingConfig",
+    "HuggingFaceHubConfig",
+    "LastCheckpointCallbackConfig",
+    "LearningRateMonitorConfig",
+    "LogEpochCallbackConfig",
+    "MetricConfig",
+    "NormLoggingCallbackConfig",
+    "OnExceptionCheckpointCallbackConfig",
+    "PluginConfigBase",
+    "RLPSanityChecksCallbackConfig",
+    "SanityCheckingConfig",
+    "SharedParametersCallbackConfig",
+    "StrategyConfigBase",
+    "TensorboardLoggerConfig",
+    "TimeCheckpointCallbackConfig",
+    "TrainerConfig",
+    "WandbLoggerConfig",
+    "_config",
+    "trainer",
+]

nshtrainer/configs/trainer/_config/__init__.py

@@ -2,16 +2,13 @@ from __future__ import annotations
 
 __codegen__ = True
 
+from nshtrainer.trainer._config import AcceleratorConfigBase as AcceleratorConfigBase
 from nshtrainer.trainer._config import ActSaveLoggerConfig as ActSaveLoggerConfig
 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 (
@@ -31,7 +28,6 @@ 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,
@@ -39,7 +35,7 @@ from nshtrainer.trainer._config import (
 from nshtrainer.trainer._config import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
-from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
+from nshtrainer.trainer._config import PluginConfigBase as PluginConfigBase
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
@@ -47,8 +43,43 @@ from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingCon
 from nshtrainer.trainer._config import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
+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
+
+__all__ = [
+    "AcceleratorConfigBase",
+    "ActSaveLoggerConfig",
+    "BaseLoggerConfig",
+    "BestCheckpointCallbackConfig",
+    "CSVLoggerConfig",
+    "CallbackConfigBase",
+    "CheckpointSavingConfig",
+    "DebugFlagCallbackConfig",
+    "DirectoryConfig",
+    "EarlyStoppingCallbackConfig",
+    "EnvironmentConfig",
+    "GradientClippingConfig",
+    "HuggingFaceHubConfig",
+    "LastCheckpointCallbackConfig",
+    "LearningRateMonitorConfig",
+    "LogEpochCallbackConfig",
+    "MetricConfig",
+    "NormLoggingCallbackConfig",
+    "OnExceptionCheckpointCallbackConfig",
+    "PluginConfigBase",
+    "RLPSanityChecksCallbackConfig",
+    "SanityCheckingConfig",
+    "SharedParametersCallbackConfig",
+    "StrategyConfigBase",
+    "TensorboardLoggerConfig",
+    "TimeCheckpointCallbackConfig",
+    "TrainerConfig",
+    "WandbLoggerConfig",
+]

nshtrainer/configs/trainer/trainer/__init__.py

@@ -2,5 +2,14 @@ from __future__ import annotations
 
 __codegen__ = True
 
+from nshtrainer.trainer.trainer import AcceleratorConfigBase as AcceleratorConfigBase
 from nshtrainer.trainer.trainer import EnvironmentConfig as EnvironmentConfig
+from nshtrainer.trainer.trainer import StrategyConfigBase as StrategyConfigBase
 from nshtrainer.trainer.trainer import TrainerConfig as TrainerConfig
+
+__all__ = [
+    "AcceleratorConfigBase",
+    "EnvironmentConfig",
+    "StrategyConfigBase",
+    "TrainerConfig",
+]

nshtrainer/configs/util/__init__.py

@@ -32,9 +32,27 @@ 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
 
 from . import _environment_info as _environment_info
 from . import config as config
+
+__all__ = [
+    "DTypeConfig",
+    "EnvironmentCUDAConfig",
+    "EnvironmentClassInformationConfig",
+    "EnvironmentConfig",
+    "EnvironmentGPUConfig",
+    "EnvironmentHardwareConfig",
+    "EnvironmentLSFInformationConfig",
+    "EnvironmentLinuxEnvironmentConfig",
+    "EnvironmentPackageConfig",
+    "EnvironmentSLURMInformationConfig",
+    "EnvironmentSnapshotConfig",
+    "EpochsConfig",
+    "GitRepositoryConfig",
+    "StepsConfig",
+    "_environment_info",
+    "config",
+]