Module optim 

Optimizers, learning-rate schedules and gradient clipping, with MLX’s formulas (mlx.optimizers).

An optimizer updates the trainable leaves of a model from a gradient tree of the model’s type (as nuvec::value_and_grad returns it): opt.update(&mut model, &grads). Its state is a tree (step, learning_rate, and the slots by parameter path), so an optimizer is part of a compiled training step’s state and of a checkpoint (nuvec::save(path, &(model, opt))):

• state is keyed by parameter path: freezing, unfreezing or replacing parameters never moves state between them, and the state of a parameter that is not updated is kept;
• the step counter and the learning rate are 0-D arrays (as MLX’s optimizer.state["learning_rate"]); a schedule writes the rate from the step at every update (evaluated before the step is counted, in float32), so a compiled step sees every new step and rate;
• hyperparameters are static fields of the tree: changing one retraces a compiled step.

Modules

schedules

Learning-rate schedules, as MLX’s mlx.optimizers schedulers: float32 array functions of the step.

Structs

Adam

Adam, as MLX’s optimizers.Adam (no weight decay; bias correction off unless enabled, as in MLX). AdamW adds decoupled weight decay.

AdamSlot

Adam’s (and AdamW’s) state: the first and second moment estimates.

AdamW

AdamW, as MLX’s optimizers.AdamW: Adam with decoupled weight decay (MLX’s default 0.01).

Lion

Lion, as MLX’s optimizers.Lion: the sign of an interpolated momentum, with optional decoupled weight decay.

MomentumSlot

Lion’s (and Muon’s) state: one momentum.

MultiOptimizer

Parameter groups, as MLX’s MultiOptimizer: each group’s filter selects parameters by path (the first matching group wins) and its optimizer updates them; the default optimizer takes the rest. Every group steps once per update. Its state is a tree: groups.0.…, …, default.….

Muon

Muon, as MLX’s optimizers.Muon: momentum (Nesterov by default) whose update is orthogonalised by Newton–Schulz iterations for parameters of two or more dimensions.

OptimizerState

An optimizer’s state: the step counter and the learning rate as 0-D arrays, and one slot (momenta, moment estimates) per parameter path.

Sgd

Stochastic gradient descent with optional momentum, dampening, Nesterov momentum and weight decay, as MLX’s optimizers.SGD.

SgdSlot

SGD’s state: the momentum buffer.

Enums

LearningRate

A learning rate: a constant or a function of the step (a 0-D array, from 0) returning a 0-D array, so a compiled step sees every new rate.

Traits

Optimizer

Updates a model’s trainable leaves from their gradients, as MLX’s Optimizer.

Functions

clip_grad_norm

grads scaled so that their joint L2 norm is at most max_norm, and that norm before scaling, as MLX’s optimizers.clip_grad_norm (float16 norms in float32). A negative max_norm gives failed arrays.