Struct Lion 

pub struct Lion {
    pub lr: f32,
    pub beta1: f32,
    pub beta2: f32,
    pub weight_decay: f32,
    /* private fields */
}

Optimizers for Func::train_step (re-exported from rlx_optim). Available with the optim feature. EvoLved sign-momentum optimizer.

Per-tensor state: one f32 buffer (half of Adam’s footprint).

Fields

lr: f32

Learning rate. Critical: typically 3–10× smaller than the AdamW LR you’d use on the same model (because the update has unit ‖sign(·)‖ per coordinate).

beta1: f32

Interpolation coefficient for the update direction (β₁ in Chen et al.). Default 0.9.

beta2: f32

EMA coefficient for the carried momentum (β₂). Default 0.99.

weight_decay: f32

Decoupled weight-decay coefficient λ. Tune ~3–10× higher than the AdamW λ you’d pair with the same model. Default 0.0.

Implementations

impl Lion

pub fn new(lr: f32) -> Lion

Construct with (β₁, β₂, λ) = (0.9, 0.99, 0.0).

pub fn with_betas(self, b1: f32, b2: f32) -> Lion

Override (β₁, β₂). They serve different roles — see the struct-level docs.

pub fn with_weight_decay(self, wd: f32) -> Lion

Override the decoupled-decay coefficient.

Trait Implementations

impl Clone for Lion

fn clone(&self) -> Lion

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for Lion

fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error>

Formats the value using the given formatter.

impl Optimizer for Lion

fn set_lr(&mut self, lr: f32)

Set the base learning rate (for LR schedules / warmup). Default is a no-op for algorithms without a scalar lr (e.g. Adafactor’s relative step size); every algorithm in this crate that has an lr field overrides this to update it.

fn step( &mut self, name: &str, _shape: &[usize], param: &mut [f32], grad: &[f32], )

fn step_batch(&mut self, items: &mut [OptItem<'_>])

Batched step over ALL parameters in one call. Default: sequential step per item — bit-identical to the per-parameter loop. Optimizers whose parameter groups are independent (e.g. Muon on the 2-D weight matrices vs AdamW on the embeddings/biases/norms) can override this to run the groups on separate threads; because the groups touch disjoint parameters and disjoint optimizer state, the result is bit-for-bit the same as the serial loop — only the wall-clock (the CPU-side optimizer bubble) shrinks toward max(group_times) instead of their sum.

fn end_iteration(&mut self)

Advance the global step counter. Most algorithms increment per call to [step], so most implementations leave this a no-op.

fn lr_scale(&self, _name: &str) -> f32

Per-tensor multiplier on the effective learning rate. Default is 1.0 for every name. Override when wrapping this crate to support per-name LR schedules (e.g. embedding-vs-attention splits, or the Gaussian-splat attribute-typed LR setup). The CPU impls in this crate currently honor this only when the caller passes a pre-scaled lr for the relevant call — backends are encouraged to consult it inside their fused kernel.