Module neuralops.trainer.registry

Classes

class TrainerRegistry

Expand source code

class TrainerRegistry:
    """Registry for trainer classes."""

    _registry: dict[str, type[BaseTrainer]] = {}

    @classmethod
    def register(cls, name: str) -> Any:
        """Register."""

        def decorator(klass: type[BaseTrainer]) -> type[BaseTrainer]:
            """Decorator."""
            if not issubclass(klass, BaseTrainer):
                raise TypeError(f"{klass.__name__} must inherit from BaseTrainer")
            cls._registry[name] = klass
            logger.info(f"Registered trainer: {name} -> {klass.__name__}")
            return klass

        return decorator

    @classmethod
    def get(cls, name: str) -> type[BaseTrainer]:
        """Get."""
        if name not in cls._registry:
            raise ValueError(f"Unknown trainer '{name}'. Available: {list(cls._registry.keys())}")
        return cls._registry[name]

    @classmethod
    def list(cls) -> list:
        """List."""
        return list(cls._registry.keys())

    @classmethod
    def create(cls, name: str, **kwargs) -> BaseTrainer:
        """Create."""
        return cls.get(name)(**kwargs)

Registry for trainer classes.

Static methods

def create(name: str, **kwargs) ‑> BaseTrainer

Create.

def get(name: str) ‑> type[BaseTrainer]

Get.

def list() ‑> TrainerRegistry.list() at 0x12f024bf0>)>

List.

def register(name: str) ‑> Any

Register.