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]Auto Trait Implementations§
impl<'a> !UnwindSafe for OptItem<'a>
impl<'a> Freeze for OptItem<'a>
impl<'a> RefUnwindSafe for OptItem<'a>
impl<'a> Send for OptItem<'a>
impl<'a> Sync for OptItem<'a>
impl<'a> Unpin for OptItem<'a>
impl<'a> UnsafeUnpin for OptItem<'a>
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more