eformer.optimizers._base

class eformer.optimizers._base.OptimizerBuilder(config: Any)

Bases: ABC

Abstract base class for optimizer builders.

Optimizer builders encapsulate the configuration and construction logic for creating optax GradientTransformation objects.

config

Optimizer-specific configuration object.

Type

Any

build()

Creates the base optimizer transformation.

validate()

Optional validation hook called before building.

abstract build(scheduler: Callable[[Union[Array, ndarray, bool, number, float, int]], Union[Array, ndarray, bool, number, float, int]]) → GradientTransformation

Build the base optimizer transformation.

Parameters

scheduler – Learning rate schedule to use.

Returns

The optimizer transformation.

Return type

optax.GradientTransformation

config: Any

validate() → None

Optional validation hook called before building the optimizer.

Raises

ValueError – If the configuration is invalid.

class eformer.optimizers._base.SchedulerBuilder(config: SchedulerConfig)

Bases: ABC

Abstract base class for scheduler builders.

Scheduler builders encapsulate the configuration and construction logic for creating optax Schedule objects.

config

Scheduler configuration object.

Type

eformer.optimizers._config.SchedulerConfig

build()

Creates the learning rate schedule.

abstract build() → Callable[[Union[Array, ndarray, bool, number, float, int]], Union[Array, ndarray, bool, number, float, int]]

Build the learning rate schedule.

Returns

The learning rate schedule.

Return type

optax.Schedule

config: SchedulerConfig

eformer.optimizers._base.register_optimizer(name: str) → Callable[[type[eformer.optimizers._base.OptimizerBuilder]], type[eformer.optimizers._base.OptimizerBuilder]]

Decorator to register an optimizer builder class.

Parameters

name – Name to register the optimizer under.

Returns

Decorator function that registers the class.

Example

@register_optimizer(“adamw”) @dataclass class AdamWOptimizer(OptimizerBuilder):

config: AdamWConfig

def build(self, scheduler):

return optax.adamw(learning_rate=scheduler, …)

eformer.optimizers._base.register_scheduler(name: str) → Callable[[type[eformer.optimizers._base.SchedulerBuilder]], type[eformer.optimizers._base.SchedulerBuilder]]

Decorator to register a scheduler builder class.

Parameters

name – Name to register the scheduler under.

Returns

Decorator function that registers the class.

Example

@register_scheduler(“cosine”) @dataclass class CosineSchedulerBuilder(SchedulerBuilder):

config: SchedulerConfig

def build(self):

return optax.cosine_decay_schedule(…)