Abstract base class owning trainable parameters and the training step.
Optimizer keeps a flat list of TensorImpl parameters and runs
one update per training iteration. Derived classes inject their
specific update mathematics via update_one and lazily allocate
per-parameter state (momentum buffers, moment estimates, parameter
averages, ...) via init_state_slot. The base class handles the
cross-cutting bookkeeping: skipping parameters without gradients,
allocating state exactly once per slot, and bumping each parameter's
version counter so autograd notices the in-place mutation.
State storage layout
Per-parameter state is keyed by an integer slot_idx that equals
the position of the parameter in params_. Derived classes hold
one std::vector<Storage> (or several) parallel to params_;
entry i is allocated on the first call to update_one for
slot i. The parallel state_initialized_ vector tracks
allocation so state is allocated exactly once.
Step contract
step() is called by the Python wrapper once per training iteration,
after the backward pass has populated param->grad_storage().
It performs three actions per slot:
- Skip the slot if the parameter pointer is null or its gradient is absent (matching the reference framework's "no grad = no update" convention).
- Lazily allocate state via
init_state_slotif this is the slot's first observed gradient. - Apply the update via
update_oneand bump the parameter's version counter viaTensorImpl::bump_versionso any autograd node that captured the parameter before the update sees the change.
zero_grad() clears every parameter's gradient storage so the next
backward pass starts from a fresh zero buffer.
Copy semantics
Copy is deleted because optimizer state owns raw heap buffers (CPU)
and MLX arrays (GPU) whose semantics under shallow copy would be
ambiguous. Optimizers are owned by the Python wrapper through a
std::shared_ptr and accessed via pybind11 reference bindings.
Serialisation
state_dict_id() returns a short version-tagged identifier (e.g.
"sgd_v1") baked into checkpoints so the Python loader can refuse
mismatched versions. state_buffers() / load_state_buffers()
expose the per-parameter mutable state as named TensorImpl lists
for torch.save-style checkpointing — entries are clones, so the
caller is free to keep them after the optimizer is destroyed.
Subclass responsibilities
Concrete optimizers must override update_one, init_state_slot,
set_lr / lr, and state_dict_id. Optimizers with global
step state (Adam, NAdam, ...) must additionally override
step_count / set_step_count. Optimizers with Python-visible
per-parameter state must override state_buffers /
load_state_buffers.
See Also
SGD, ASGD : plain and averaged stochastic gradient descent. Adam, AdamW, NAdam, RAdam, Adamax, Adafactor : adaptive moment family. Adagrad, Adadelta, RMSProp : adaptive learning-rate family.
Attributes
params_std::vector<std::shared_ptr<TensorImpl>>slot_idx in the protected hooks. May contain null entries; nulls are silently skipped by step().state_initialized_std::vector<bool>params_. Marks whether init_state_slot has run for each slot; grown lazily by step() to match params_.size() on first use.Constructors
1Construct an optimizer from a flat list of trainable parameters.
Parameters
paramsstd::vector<std::shared_ptr<TensorImpl>>step). The caller's vector is moved into the optimizer; subsequent mutations of the caller's copy do not affect this optimizer.Notes
state_initialized_ is sized lazily on the first step()
rather than here, which lets the Python wrapper extend params_
post-construction in the rare param-group reconfiguration case.
Destructor
1Virtual destructor — ensures derived Adam / SGD / RMSProp destructors run when held by a base-class pointer. Defaulted; the per-parameter state buffers are released by their own std::vector members in the derived classes, and params_ releases its shared_ptr refs automatically.
Operators
1Virtual methods
6Allocate and zero-initialise per-parameter state for one slot.
Parameters
slot_idxstd::size_tparams_ whose state should be materialised.paramconst std::shared_ptr<TensorImpl>&Notes
Called exactly once per slot — on the first step() for which
that slot has a gradient. Implementations should size their
std::vector<Storage> state members lazily here rather than in
the constructor so that disabled features (e.g. momentum == 0
in SGD) skip the allocation entirely.
Restore the per-parameter state in place from a snapshot.
See Also
state_buffers : the inverse operation.
Parameters
bufsconst std::vector<NamedBuffers>&state_buffers for this optimizer: same buffer names, same length, same per-slot shapes / dtypes. Mismatches raise.Raises
:runtime_errorNotes
After this returns, every state slot is treated as initialised so
step() will not re-zero it on the next call. Implementations
typically delegate the byte copy to overwrite_state_storage.
Current learning rate.
Returns
doubleThe value last passed to set_lr (or the constructor).
Set the learning rate (called by LR schedulers between steps).
Parameters
lrdoubleOverride the global step counter (used by load_state_dict).
Parameters
countstd::int64_tApply the optimizer's update rule for a single parameter.
Parameters
slot_idxstd::size_tparams_. Used to look up the pre-allocated state buffers (momentum, second moment, parameter average, ...).paramstd::shared_ptr<TensorImpl>&gradconst Storage&Notes
Implementations must mutate param's storage in place — the
base step() is responsible for bumping the version counter
after this returns. Both CPU and GPU paths are typically
dispatched from inside this method based on param->device().
Methods
6Number of parameter slots managed by this optimizer.
Returns
std::size_tparams_.size() — includes null entries.
Snapshot the per-parameter mutable state for checkpointing.
See Also
load_state_buffers : the inverse operation.
Returns
std::vector<NamedBuffers>A list of (name, tensors) pairs. name identifies the buffer kind (e.g. "momentum_buffer"); tensors runs parallel to params_ with one entry per parameter slot (or a null entry if no state was allocated for that slot). Entries are clones — callers may keep them after the optimizer is destroyed. Optimizers with no Python-visible state return an empty vector.
Notes
Implementations typically call clone_state_storage for each
active slot. Slots whose state was never allocated (e.g. SGD
with momentum == 0) should contribute a null pointer rather
than a synthetic zero buffer.
Short version-tagged identifier baked into checkpoints.
Returns
std::stringA stable short tag like "sgd_v1" or "adam_v1" used by the Python loader to detect mismatched optimizer versions in load_state_dict.
Run one optimizer update over all registered parameters.
For each non-null parameter that has an accumulated gradient,
step() lazily initialises its state slot, invokes the derived
class's update_one to apply the update rule, then bumps the
parameter's autograd version counter. Parameters without a
gradient are skipped — this matches the reference framework's
"no grad = no update" semantics and lets a single optimizer
safely cover partially active sub-modules (e.g. frozen layers).
Notes
Called once per training iteration, after loss.backward().
The accompanying Python wrapper takes care of forcing the
resulting MLX graph to eval() only when requested via
AUTO_EVAL_AFTER_STEP; by default step() is itself lazy
so that forward → backward → step can fuse into one MLX
submission.
Global step counter used by some optimizers for bias correction.
Returns
std::int64_tCurrent step count (Adam bias correction, NAdam momentum schedule, ...). Optimizers without a global counter return 0.
Zero every parameter's gradient storage.
Calls TensorImpl::zero_grad on each non-null parameter.
Should be called before each forward/backward pass so that the
next backward accumulates into a clean buffer rather than into
the previous step's gradient.