overlay/hydra/__init__.py

"""HYDRA training package.

Thin facade re-exporting the public API used by train.py, the test suite,
and external research scripts. Imports are lazy where possible to keep
`import hydra` cheap (prepare.py and mamba-ssm are the heavy deps).
"""

from hydra.config import PostSemClawConfig
from hydra.engram import GPUEngram
from hydra.optimizer import MuonAdamW, adamw_step_fused, muon_step_fused

# Heavy imports are resolved lazily so `import hydra` and `import hydra.hyena_block`
# keep working in local CPU/test environments that do not have the container-only
# mamba-ssm wheel stack installed.
def __getattr__(name: str):
    if name == "PostSemClawModel":
        from hydra.model import PostSemClawModel as _model
        return _model
    if name == "norm":
        from hydra.model import norm as _norm
        return _norm
    if name == "config_from_dict":
        from hydra.training import config_from_dict as _cfd
        return _cfd
    raise AttributeError(name)


__all__ = [
    "PostSemClawConfig",
    "GPUEngram",
    "PostSemClawModel",
    "norm",
    "MuonAdamW",
    "adamw_step_fused",
    "muon_step_fused",
    "config_from_dict",
]