pub struct OptItem<'a> {
pub name: &'a str,
pub shape: &'a [usize],
pub param: &'a mut [f32],
pub grad: &'a [f32],
}Expand description
Common parameter-update interface.
name keys the per-parameter state (moments, preconditioners),
shape is the parameter’s logical shape (used by matrix-aware
algorithms like Adafactor / SOAP / Muon — ignored by elementwise
ones), param is updated in place from grad. grad is treated
as read-only; callers that need gradient clipping should pre-scale
it (see global_grad_clip_scale).
§Implementing for a backend
Every algorithm in this crate provides a CPU reference impl. A
backend (e.g. rlx-metal, rlx-cuda) is free to write its own
fused step kernel and impl Optimizer for a wrapper struct that
owns device buffers — the trait places no requirement on where
the state lives, only on the entry-point signature. The
rlx-metal::splat_adam kernel is the canonical example of a
backend that bypasses this crate entirely; you can wrap it with a
5-line impl Optimizer if you want a uniform interface from a
generic trainer.
§Per-tensor learning rate
For optimizers that don’t need per-tensor LR variation (most
transformer pre-training), set lr_scale to
return 1.0 (the default). For domain-specific use cases — e.g.
3D Gaussian splatting, where different attributes need wildly
different step sizes — override lr_scale to
multiply the base lr by a per-name factor. The provided method
on the trait does NOT scale automatically; algorithms are free to
consult it via Optimizer::lr_scale inside their step.
One parameter’s data for a batched optimizer step (Optimizer::step_batch):
its name, static shape, mutable data slice (owned by the caller), and gradient.
Fields§
§name: &'a str§shape: &'a [usize]§param: &'a mut [f32]§grad: &'a [f32]