Module optim 

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

Every optimizer updates the trainable parameters of a Module from gradients in the order of Module::trainable_parameters (as crate::value_and_grad returns them). The learning rate is a constant or a schedules function of the step, evaluated before the step is counted and kept in float32, as in MLX.

Modules

schedules

Learning-rate schedules, as MLX’s mlx.optimizers schedulers (computed in float32).

Structs

Adam

Adam, as MLX’s optimizers.Adam (bias correction off unless enabled, as in MLX), and AdamW (decoupled weight decay) with Adam::weight_decay.

Lion

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

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.

Sgd

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

Enums

LearningRate

A learning rate: a constant or a function of the step (from 0).

Traits

Optimizer

Updates a model’s trainable parameters from their gradients.

Functions

clip_grad_norm

gradients 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).