ml.optimizers.adan

Wrapper around the PyTorch Adan optimizer.

class ml.optimizers.adan.Adan(params: Iterable[Parameter] | Iterable[dict[str, Any]], lr: float = 0.001, betas: tuple[float, float, float] = (0.1, 0.1, 0.001), eps: float = 1e-08, weight_decay: float = 0.0)

Bases: Optimizer

step(closure: Callable[[], float] | None = None) → float | None

Performs a single optimization step (parameter update).

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

Note

Unless otherwise specified, this function should not modify the .grad field of the parameters.

class ml.optimizers.adan.AdanOptimizerConfig(name: str = '???', lr: float = 0.001, betas: tuple[float, float, float] = (0.1, 0.1, 0.001), eps: float = 0.0001, weight_decay: float = 1e-05, default_decay: bool = True)

Bases: BaseOptimizerConfig

lr: float = 0.001

betas: tuple[float, float, float] = (0.1, 0.1, 0.001)

eps: float = 0.0001

weight_decay: float = 1e-05

default_decay: bool = True

class ml.optimizers.adan.AdanOptimizer(config: BaseConfigT)

Bases: BaseOptimizer[AdanOptimizerConfig, Adan]

get(model: Module) → Adan

Given a base module, returns an optimizer.

Parameters:

model – The model to get an optimizer for

Returns:

The constructed optimizer