nshtrainer 0.41.1__py3-none-any.whl → 0.42.0__py3-none-any.whl

nshtrainer/config/loggers/wandb/__init__.py

@@ -1,16 +1,36 @@
-# fmt: off
-# ruff: noqa
-# type: ignore
-
 __codegen__ = True
 
-# Config classes
-from nshtrainer.loggers.wandb import BaseLoggerConfig as BaseLoggerConfig
-from nshtrainer.loggers.wandb import WandbLoggerConfig as WandbLoggerConfig
-from nshtrainer.loggers.wandb import WandbUploadCodeConfig as WandbUploadCodeConfig
-from nshtrainer.loggers.wandb import CallbackConfigBase as CallbackConfigBase
-from nshtrainer.loggers.wandb import WandbWatchConfig as WandbWatchConfig
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.loggers.wandb import BaseLoggerConfig as BaseLoggerConfig
+    from nshtrainer.loggers.wandb import CallbackConfigBase as CallbackConfigBase
+    from nshtrainer.loggers.wandb import WandbLoggerConfig as WandbLoggerConfig
+    from nshtrainer.loggers.wandb import WandbUploadCodeConfig as WandbUploadCodeConfig
+    from nshtrainer.loggers.wandb import WandbWatchConfig as WandbWatchConfig
+else:
+
+    def __getattr__(name):
+        import importlib
 
-# Type aliases
+        if name in globals():
+            return globals()[name]
+        if name == "CallbackConfigBase":
+            return importlib.import_module(
+                "nshtrainer.loggers.wandb"
+            ).CallbackConfigBase
+        if name == "WandbLoggerConfig":
+            return importlib.import_module("nshtrainer.loggers.wandb").WandbLoggerConfig
+        if name == "WandbUploadCodeConfig":
+            return importlib.import_module(
+                "nshtrainer.loggers.wandb"
+            ).WandbUploadCodeConfig
+        if name == "WandbWatchConfig":
+            return importlib.import_module("nshtrainer.loggers.wandb").WandbWatchConfig
+        if name == "BaseLoggerConfig":
+            return importlib.import_module("nshtrainer.loggers.wandb").BaseLoggerConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/lr_scheduler/__init__.py

@@ -1,18 +1,55 @@
-# fmt: off
-# ruff: noqa
-# type: ignore
-
 __codegen__ = True
 
-# Config classes
-from nshtrainer.lr_scheduler import LRSchedulerConfigBase as LRSchedulerConfigBase
-from nshtrainer.lr_scheduler import LinearWarmupCosineDecayLRSchedulerConfig as LinearWarmupCosineDecayLRSchedulerConfig
-from nshtrainer.lr_scheduler import ReduceLROnPlateauConfig as ReduceLROnPlateauConfig
-from nshtrainer.lr_scheduler.reduce_lr_on_plateau import MetricConfig as MetricConfig
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    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,
+    )
+else:
+
+    def __getattr__(name):
+        import importlib
+
+        if name in globals():
+            return globals()[name]
+        if name == "LRSchedulerConfigBase":
+            return importlib.import_module(
+                "nshtrainer.lr_scheduler"
+            ).LRSchedulerConfigBase
+        if name == "LinearWarmupCosineDecayLRSchedulerConfig":
+            return importlib.import_module(
+                "nshtrainer.lr_scheduler"
+            ).LinearWarmupCosineDecayLRSchedulerConfig
+        if name == "MetricConfig":
+            return importlib.import_module(
+                "nshtrainer.lr_scheduler.reduce_lr_on_plateau"
+            ).MetricConfig
+        if name == "ReduceLROnPlateauConfig":
+            return importlib.import_module(
+                "nshtrainer.lr_scheduler"
+            ).ReduceLROnPlateauConfig
+        if name == "DurationConfig":
+            return importlib.import_module(
+                "nshtrainer.lr_scheduler.linear_warmup_cosine"
+            ).DurationConfig
+        if name == "LRSchedulerConfig":
+            return importlib.import_module("nshtrainer.lr_scheduler").LRSchedulerConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
-# Type aliases
-from nshtrainer.lr_scheduler.linear_warmup_cosine import DurationConfig as DurationConfig
-from nshtrainer.lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
 
 # Submodule exports
 from . import _base as _base

nshtrainer/config/lr_scheduler/_base/__init__.py

@@ -1,12 +1,24 @@
-# fmt: off
-# ruff: noqa
-# type: ignore
-
 __codegen__ = True
 
-# Config classes
-from nshtrainer.lr_scheduler._base import LRSchedulerConfigBase as LRSchedulerConfigBase
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.lr_scheduler._base import (
+        LRSchedulerConfigBase as LRSchedulerConfigBase,
+    )
+else:
+
+    def __getattr__(name):
+        import importlib
 
-# Type aliases
+        if name in globals():
+            return globals()[name]
+        if name == "LRSchedulerConfigBase":
+            return importlib.import_module(
+                "nshtrainer.lr_scheduler._base"
+            ).LRSchedulerConfigBase
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/lr_scheduler/linear_warmup_cosine/__init__.py

@@ -1,14 +1,38 @@
-# fmt: off
-# ruff: noqa
-# type: ignore
-
 __codegen__ = True
 
-# Config classes
-from nshtrainer.lr_scheduler.linear_warmup_cosine import LinearWarmupCosineDecayLRSchedulerConfig as LinearWarmupCosineDecayLRSchedulerConfig
-from nshtrainer.lr_scheduler.linear_warmup_cosine import LRSchedulerConfigBase as LRSchedulerConfigBase
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    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,
+    )
+else:
+
+    def __getattr__(name):
+        import importlib
 
-# Type aliases
-from nshtrainer.lr_scheduler.linear_warmup_cosine import DurationConfig as DurationConfig
+        if name in globals():
+            return globals()[name]
+        if name == "LinearWarmupCosineDecayLRSchedulerConfig":
+            return importlib.import_module(
+                "nshtrainer.lr_scheduler.linear_warmup_cosine"
+            ).LinearWarmupCosineDecayLRSchedulerConfig
+        if name == "LRSchedulerConfigBase":
+            return importlib.import_module(
+                "nshtrainer.lr_scheduler.linear_warmup_cosine"
+            ).LRSchedulerConfigBase
+        if name == "DurationConfig":
+            return importlib.import_module(
+                "nshtrainer.lr_scheduler.linear_warmup_cosine"
+            ).DurationConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/lr_scheduler/reduce_lr_on_plateau/__init__.py

@@ -1,14 +1,38 @@
-# fmt: off
-# ruff: noqa
-# type: ignore
-
 __codegen__ = True
 
-# Config classes
-from nshtrainer.lr_scheduler.reduce_lr_on_plateau import ReduceLROnPlateauConfig as ReduceLROnPlateauConfig
-from nshtrainer.lr_scheduler.reduce_lr_on_plateau import LRSchedulerConfigBase as LRSchedulerConfigBase
-from nshtrainer.lr_scheduler.reduce_lr_on_plateau import MetricConfig as MetricConfig
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.lr_scheduler.reduce_lr_on_plateau import (
+        LRSchedulerConfigBase as LRSchedulerConfigBase,
+    )
+    from nshtrainer.lr_scheduler.reduce_lr_on_plateau import (
+        MetricConfig as MetricConfig,
+    )
+    from nshtrainer.lr_scheduler.reduce_lr_on_plateau import (
+        ReduceLROnPlateauConfig as ReduceLROnPlateauConfig,
+    )
+else:
+
+    def __getattr__(name):
+        import importlib
 
-# Type aliases
+        if name in globals():
+            return globals()[name]
+        if name == "MetricConfig":
+            return importlib.import_module(
+                "nshtrainer.lr_scheduler.reduce_lr_on_plateau"
+            ).MetricConfig
+        if name == "ReduceLROnPlateauConfig":
+            return importlib.import_module(
+                "nshtrainer.lr_scheduler.reduce_lr_on_plateau"
+            ).ReduceLROnPlateauConfig
+        if name == "LRSchedulerConfigBase":
+            return importlib.import_module(
+                "nshtrainer.lr_scheduler.reduce_lr_on_plateau"
+            ).LRSchedulerConfigBase
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/metrics/__init__.py

@@ -1,13 +1,22 @@
-# fmt: off
-# ruff: noqa
-# type: ignore
-
 __codegen__ = True
 
-# Config classes
-from nshtrainer.metrics import MetricConfig as MetricConfig
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.metrics import MetricConfig as MetricConfig
+else:
+
+    def __getattr__(name):
+        import importlib
+
+        if name in globals():
+            return globals()[name]
+        if name == "MetricConfig":
+            return importlib.import_module("nshtrainer.metrics").MetricConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
-# Type aliases
 
 # Submodule exports
 from . import _config as _config

nshtrainer/config/metrics/_config/__init__.py

@@ -1,12 +1,20 @@
-# fmt: off
-# ruff: noqa
-# type: ignore
-
 __codegen__ = True
 
-# Config classes
-from nshtrainer.metrics._config import MetricConfig as MetricConfig
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.metrics._config import MetricConfig as MetricConfig
+else:
+
+    def __getattr__(name):
+        import importlib
 
-# Type aliases
+        if name in globals():
+            return globals()[name]
+        if name == "MetricConfig":
+            return importlib.import_module("nshtrainer.metrics._config").MetricConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/model/__init__.py

@@ -1,18 +1,37 @@
-# fmt: off
-# ruff: noqa
-# type: ignore
-
 __codegen__ = True
 
-# Config classes
-from nshtrainer.model import MetricConfig as MetricConfig
-from nshtrainer.model import BaseConfig as BaseConfig
-from nshtrainer.model import DirectoryConfig as DirectoryConfig
-from nshtrainer.model import TrainerConfig as TrainerConfig
-from nshtrainer.model.base import EnvironmentConfig as EnvironmentConfig
-from nshtrainer.model.config import CallbackConfigBase as CallbackConfigBase
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.model import BaseConfig as BaseConfig
+    from nshtrainer.model import DirectoryConfig as DirectoryConfig
+    from nshtrainer.model import MetricConfig as MetricConfig
+    from nshtrainer.model import TrainerConfig as TrainerConfig
+    from nshtrainer.model.base import EnvironmentConfig as EnvironmentConfig
+    from nshtrainer.model.config import CallbackConfigBase as CallbackConfigBase
+else:
+
+    def __getattr__(name):
+        import importlib
+
+        if name in globals():
+            return globals()[name]
+        if name == "CallbackConfigBase":
+            return importlib.import_module("nshtrainer.model.config").CallbackConfigBase
+        if name == "TrainerConfig":
+            return importlib.import_module("nshtrainer.model").TrainerConfig
+        if name == "EnvironmentConfig":
+            return importlib.import_module("nshtrainer.model.base").EnvironmentConfig
+        if name == "BaseConfig":
+            return importlib.import_module("nshtrainer.model").BaseConfig
+        if name == "MetricConfig":
+            return importlib.import_module("nshtrainer.model").MetricConfig
+        if name == "DirectoryConfig":
+            return importlib.import_module("nshtrainer.model").DirectoryConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
-# Type aliases
 
 # Submodule exports
 from . import base as base

nshtrainer/config/model/base/__init__.py

@@ -1,13 +1,23 @@
-# fmt: off
-# ruff: noqa
-# type: ignore
-
 __codegen__ = True
 
-# Config classes
-from nshtrainer.model.base import EnvironmentConfig as EnvironmentConfig
-from nshtrainer.model.base import BaseConfig as BaseConfig
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.model.base import BaseConfig as BaseConfig
+    from nshtrainer.model.base import EnvironmentConfig as EnvironmentConfig
+else:
+
+    def __getattr__(name):
+        import importlib
 
-# Type aliases
+        if name in globals():
+            return globals()[name]
+        if name == "EnvironmentConfig":
+            return importlib.import_module("nshtrainer.model.base").EnvironmentConfig
+        if name == "BaseConfig":
+            return importlib.import_module("nshtrainer.model.base").BaseConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/model/config/__init__.py

@@ -1,17 +1,35 @@
-# fmt: off
-# ruff: noqa
-# type: ignore
-
 __codegen__ = True
 
-# Config classes
-from nshtrainer.model.config import MetricConfig as MetricConfig
-from nshtrainer.model.config import BaseConfig as BaseConfig
-from nshtrainer.model.config import DirectoryConfig as DirectoryConfig
-from nshtrainer.model.config import TrainerConfig as TrainerConfig
-from nshtrainer.model.config import EnvironmentConfig as EnvironmentConfig
-from nshtrainer.model.config import CallbackConfigBase as CallbackConfigBase
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.model.config import BaseConfig as BaseConfig
+    from nshtrainer.model.config import CallbackConfigBase as CallbackConfigBase
+    from nshtrainer.model.config import DirectoryConfig as DirectoryConfig
+    from nshtrainer.model.config import EnvironmentConfig as EnvironmentConfig
+    from nshtrainer.model.config import MetricConfig as MetricConfig
+    from nshtrainer.model.config import TrainerConfig as TrainerConfig
+else:
+
+    def __getattr__(name):
+        import importlib
 
-# Type aliases
+        if name in globals():
+            return globals()[name]
+        if name == "CallbackConfigBase":
+            return importlib.import_module("nshtrainer.model.config").CallbackConfigBase
+        if name == "TrainerConfig":
+            return importlib.import_module("nshtrainer.model.config").TrainerConfig
+        if name == "EnvironmentConfig":
+            return importlib.import_module("nshtrainer.model.config").EnvironmentConfig
+        if name == "BaseConfig":
+            return importlib.import_module("nshtrainer.model.config").BaseConfig
+        if name == "MetricConfig":
+            return importlib.import_module("nshtrainer.model.config").MetricConfig
+        if name == "DirectoryConfig":
+            return importlib.import_module("nshtrainer.model.config").DirectoryConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/model/mixins/logger/__init__.py

@@ -1,12 +1,20 @@
-# fmt: off
-# ruff: noqa
-# type: ignore
-
 __codegen__ = True
 
-# Config classes
-from nshtrainer.model.mixins.logger import BaseConfig as BaseConfig
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.model.mixins.logger import BaseConfig as BaseConfig
+else:
+
+    def __getattr__(name):
+        import importlib
 
-# Type aliases
+        if name in globals():
+            return globals()[name]
+        if name == "BaseConfig":
+            return importlib.import_module("nshtrainer.model.mixins.logger").BaseConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/nn/__init__.py

@@ -1,29 +1,74 @@
-# fmt: off
-# ruff: noqa
-# type: ignore
-
 __codegen__ = True
 
-# Config classes
-from nshtrainer.nn import MLPConfig as MLPConfig
-from nshtrainer.nn import BaseNonlinearityConfig as BaseNonlinearityConfig
-from nshtrainer.nn import TanhNonlinearityConfig as TanhNonlinearityConfig
-from nshtrainer.nn import ReLUNonlinearityConfig as ReLUNonlinearityConfig
-from nshtrainer.nn import PReLUConfig as PReLUConfig
-from nshtrainer.nn import SiLUNonlinearityConfig as SiLUNonlinearityConfig
-from nshtrainer.nn import LeakyReLUNonlinearityConfig as LeakyReLUNonlinearityConfig
-from nshtrainer.nn import SoftsignNonlinearityConfig as SoftsignNonlinearityConfig
-from nshtrainer.nn import MishNonlinearityConfig as MishNonlinearityConfig
-from nshtrainer.nn import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
-from nshtrainer.nn import SoftplusNonlinearityConfig as SoftplusNonlinearityConfig
-from nshtrainer.nn.nonlinearity import SwiGLUNonlinearityConfig as SwiGLUNonlinearityConfig
-from nshtrainer.nn import ELUNonlinearityConfig as ELUNonlinearityConfig
-from nshtrainer.nn import SoftmaxNonlinearityConfig as SoftmaxNonlinearityConfig
-from nshtrainer.nn import GELUNonlinearityConfig as GELUNonlinearityConfig
-from nshtrainer.nn import SwishNonlinearityConfig as SwishNonlinearityConfig
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.nn import BaseNonlinearityConfig as BaseNonlinearityConfig
+    from nshtrainer.nn import ELUNonlinearityConfig as ELUNonlinearityConfig
+    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
+    from nshtrainer.nn import SiLUNonlinearityConfig as SiLUNonlinearityConfig
+    from nshtrainer.nn import SoftmaxNonlinearityConfig as SoftmaxNonlinearityConfig
+    from nshtrainer.nn import SoftplusNonlinearityConfig as SoftplusNonlinearityConfig
+    from nshtrainer.nn import SoftsignNonlinearityConfig as SoftsignNonlinearityConfig
+    from nshtrainer.nn import SwishNonlinearityConfig as SwishNonlinearityConfig
+    from nshtrainer.nn import TanhNonlinearityConfig as TanhNonlinearityConfig
+    from nshtrainer.nn.nonlinearity import (
+        SwiGLUNonlinearityConfig as SwiGLUNonlinearityConfig,
+    )
+else:
+
+    def __getattr__(name):
+        import importlib
+
+        if name in globals():
+            return globals()[name]
+        if name == "BaseNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").BaseNonlinearityConfig
+        if name == "MLPConfig":
+            return importlib.import_module("nshtrainer.nn").MLPConfig
+        if name == "SwiGLUNonlinearityConfig":
+            return importlib.import_module(
+                "nshtrainer.nn.nonlinearity"
+            ).SwiGLUNonlinearityConfig
+        if name == "ReLUNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").ReLUNonlinearityConfig
+        if name == "SiLUNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SiLUNonlinearityConfig
+        if name == "ELUNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").ELUNonlinearityConfig
+        if name == "GELUNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").GELUNonlinearityConfig
+        if name == "SoftplusNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SoftplusNonlinearityConfig
+        if name == "SoftsignNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SoftsignNonlinearityConfig
+        if name == "SwishNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SwishNonlinearityConfig
+        if name == "SoftmaxNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SoftmaxNonlinearityConfig
+        if name == "MishNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").MishNonlinearityConfig
+        if name == "SigmoidNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SigmoidNonlinearityConfig
+        if name == "TanhNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").TanhNonlinearityConfig
+        if name == "PReLUConfig":
+            return importlib.import_module("nshtrainer.nn").PReLUConfig
+        if name == "LeakyReLUNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").LeakyReLUNonlinearityConfig
+        if name == "NonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").NonlinearityConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
-# Type aliases
-from nshtrainer.nn import NonlinearityConfig as NonlinearityConfig
 
 # Submodule exports
 from . import mlp as mlp

nshtrainer/config/nn/mlp/__init__.py

@@ -1,14 +1,26 @@
-# fmt: off
-# ruff: noqa
-# type: ignore
-
 __codegen__ = True
 
-# Config classes
-from nshtrainer.nn.mlp import MLPConfig as MLPConfig
-from nshtrainer.nn.mlp import BaseNonlinearityConfig as BaseNonlinearityConfig
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.nn.mlp import BaseNonlinearityConfig as BaseNonlinearityConfig
+    from nshtrainer.nn.mlp import MLPConfig as MLPConfig
+    from nshtrainer.nn.mlp import NonlinearityConfig as NonlinearityConfig
+else:
+
+    def __getattr__(name):
+        import importlib
 
-# Type aliases
-from nshtrainer.nn.mlp import NonlinearityConfig as NonlinearityConfig
+        if name in globals():
+            return globals()[name]
+        if name == "BaseNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn.mlp").BaseNonlinearityConfig
+        if name == "MLPConfig":
+            return importlib.import_module("nshtrainer.nn.mlp").MLPConfig
+        if name == "NonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn.mlp").NonlinearityConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports