# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 Andrew Hendel. Part of mojolearn, https://doi.org/10.5281/zenodo.22068632
"""torch.optim-shaped optimizers of the sequence expansion lane, over float32
NumPy arrays updated in place: `opt = RMSprop([w1, w2], lr=1e-2);
opt.step([g1, g2])`. One launch over the packed flat buffer per step, the
element body `sequence/ops.mojo::op_opt` (PyTorch's update rules), the same on
the GPU binding and on the CPU host binding. `lr` is read fresh at every step,
so a schedule is `opt.lr = ...` between steps, or `opt.lr_schedule = <schedule>` (anything
with `lr_at(t)`, t one-based: `StepLR`, `ExponentialLR`, `OneCycleLR`), which
sets `lr` before every step.
Refused (sequence/NOT_IMPLEMENTED.tsv): parameter groups, maximize,
foreach/fused/capturable/differentiable, sparse gradients, float64."""
from ._optional_numpy import require_numpy
np = require_numpy('_x_sequence_optim')
import os
from . import _backend
from ._x_sequence_rnn import OPTIMIZERS, optimizer_arguments
from ._x_sequence_device import SequenceDeviceTensor, is_seq_tensor, resident_binding
#: lane gap-optimizers (2026-10-02): the state slots live on the device
#: binding's context across steps (`sequence/opt_resident.mojo`), so a step
#: moves only the parameters and gradients up and the parameters down, each
#: tensor straight into its place (no host concatenation); the same element
#: statements on the same values, the same bits. `MOJOLEARN_OPTIMIZER_RESIDENT=0`
#: keeps the per-call entries (the before arm of the A/B); the host binding
#: has no resident entries and keeps them too.
_RESIDENT_ENV = "MOJOLEARN_OPTIMIZER_RESIDENT"
def _resident_binding(b, entries):
if os.environ.get(_RESIDENT_ENV, "1") == "0":
return False
return all(callable(getattr(b, n, None)) for n in entries) # glue: the named binding entry points
_SEQ_RESIDENT = ("optimizer_resident_open", "optimizer_resident_close", "optimizer_resident_move",
"optimizer_resident_step")
_AF_RESIDENT = ("adafactor_resident_open", "optimizer_resident_close", "optimizer_resident_move",
"adafactor_resident_step")
_LAMB_RESIDENT = ("lamb_resident_open", "optimizer_resident_close", "optimizer_resident_move",
"lamb_resident_step")
#: lane gap-neural-io (2026-10-08, plan docs/plans/gaps-2026-10-08.md Section
#: 4): params given as `SequenceDeviceTensor`s (`_x_sequence_device.to_device`)
#: are updated where they live by the `_dev` steps (`sequence/opt_resident.mojo`):
#: no upload of P and G and no download of P per step (3 x 67 MB at the
#: board's 16,777,216 floats), the same launches on the same values, the same
#: bits. The state stays resident as above; grads may be device tensors or
#: host arrays (uploaded once per step, here).
_SEQ_DEV = _SEQ_RESIDENT + ("optimizer_resident_step_dev",)
_AF_DEV = _AF_RESIDENT + ("adafactor_resident_step_dev",)
_LAMB_DEV = _LAMB_RESIDENT + ("lamb_resident_step_dev",)
def _device_params(params, who, entries):
"""None when no param is a `SequenceDeviceTensor`; their binding when
every one is (on one binding that has `entries`); refused otherwise."""
dev = [is_seq_tensor(p) for p in params] # glue: per parameter tensor, not elements
if not any(dev):
return None
if not all(dev):
raise TypeError(f"{who}: params are all device tensors or all NumPy arrays, not a mix")
b = params[0].b
if any(p.b is not b for p in params): # glue: per parameter tensor, not elements
raise TypeError(f"{who}: every device param must live on one binding")
if not (resident_binding(b) and all(callable(getattr(b, n, None)) for n in entries)): # glue: the named binding entry points
raise RuntimeError(f"{who}: this binding has no device-parameter step; pass NumPy arrays")
if os.environ.get(_RESIDENT_ENV, "1") == "0":
raise RuntimeError(f"{who}: device params need the resident state ({_RESIDENT_ENV}=0 is set)")
return b
def _device_grads(b, grads, params, who):
"""The grads as device tensors of the params' shapes on binding `b`; a
host float32 array is uploaded once (here)."""
if isinstance(grads, (np.ndarray, SequenceDeviceTensor)):
grads = [grads]
if len(grads) != len(params):
raise ValueError(f"{who}: {len(grads)} grads, {len(params)} params")
out = []
for g, p in zip(grads, params): # glue: per parameter tensor, not elements
if is_seq_tensor(g):
if g.b is not b or tuple(g.shape) != tuple(p.shape):
raise ValueError(f"{who}: a device grad of shape {g.shape} for a param {p.shape} "
"(or on another binding)")
out.append(g)
continue
a = np.asarray(g)
if a.shape != tuple(p.shape):
raise ValueError(f"{who}: a grads of shape {a.shape} for a param {p.shape}")
if a.dtype != np.float32:
raise TypeError(f"{who}: grads must be float32 (got {a.dtype})")
out.append(SequenceDeviceTensor._wrap(b, a))
return out
class _ResidentState:
"""The state slots of a resident handle, mirrored on the host on access:
`state` downloads the slots the handle keeps (once per stale read), and
the next step uploads the host copies first (the caller may have
written them, `load_state_dict` does)."""
def _res_init(self, state):
self._state = state
self._res = None
self._res_b = None
self._res_used = 0
self._host_fresh = True # the host copies hold the newest values
self._host_owned = True # the next step uploads them first
# the host copies are still the zeros (and initial values) the
# constructor built: nobody has read or replaced them
self._pristine = True
@property
def state(self):
if self._res is not None and not self._host_fresh:
for k, s in enumerate(self._state): # glue: the optimizer state slots, at most three
self._res_b.optimizer_resident_move(self._res, k, s.ctypes.data, 0)
self._host_fresh = True
# a reader may write into the arrays it got: upload before the next step
self._host_owned = True
self._pristine = False
return self._state
@state.setter
def state(self, value):
self._state = value
self._host_fresh = True
self._host_owned = True
self._pristine = False
def _res_opened(self, b, ret, init_zero=True):
"""Keep the handle an open returned ([handle, used] or, from a
FAST+Apple binding (OPT_ZERO_OPEN, the default there), [handle, used, 1]: the
device slots are zero filled, so host copies that are still the
constructor's zeros need no upload)."""
self._res, self._res_b, self._res_used = int(ret[0]), b, int(ret[1])
zeroed = len(ret) > 2 and int(ret[2]) == 1
self._host_owned = not (zeroed and self._pristine and init_zero)
def _res_upload(self):
if self._host_owned:
for k, s in enumerate(self._state): # glue: the optimizer state slots, at most three
self._res_b.optimizer_resident_move(self._res, k, s.ctypes.data, 1)
self._host_owned = False
self._host_fresh = False
@property
def resident_(self):
"""True once the state lives on the device."""
return self._res is not None
def __del__(self):
res, b = getattr(self, "_res", None), getattr(self, "_res_b", None)
if res is not None and b is not None:
try:
b.optimizer_resident_close(res)
except Exception:
pass
def _grad_list(self, grads):
"""The gradients as C-contiguous float32 arrays of the params' shapes
(copied only when they are not already)."""
if isinstance(grads, np.ndarray):
grads = [grads]
if len(grads) != len(self.params):
raise ValueError(f"{type(self).__name__}: {len(grads)} grads, {len(self.params)} params")
out = []
for a, p in zip(grads, self.params): # glue: per parameter tensor, not elements
if np.shape(a) != p.shape:
raise ValueError(f"{type(self).__name__}: a grads of shape {np.shape(a)} for a param {p.shape}")
if np.asarray(a).dtype == np.float64:
raise TypeError(f"{type(self).__name__}: float64 grads are refused")
out.append(np.ascontiguousarray(a, dtype=np.float32))
return out
def _check_params(self):
for k, p in enumerate(self.params): # glue: per parameter tensor, not elements
if not isinstance(p, np.ndarray) or p.dtype != np.float32 or not p.flags.c_contiguous \
or not p.flags.writeable:
raise TypeError(f"{type(self).__name__}: params[{k}] must stay a writable C-contiguous "
"float32 NumPy array (it is updated in place)")
if sum(p.size for p in self.params) != self.n_total: # glue: sizes per parameter tensor, not elements
raise ValueError(f"{type(self).__name__}: the params were resized under the optimizer")
class _SeqOptimizer(_ResidentState):
_NAME = None
def __init__(self, params, lr, options, numeric_mode=None):
if isinstance(params, np.ndarray):
params = [params]
self.params = list(params)
self._dev_b = _device_params(self.params, type(self).__name__, _SEQ_DEV)
for k, p in enumerate(self.params): # glue: per parameter tensor, not elements
if self._dev_b is None and (not isinstance(p, np.ndarray) or p.dtype != np.float32
or not p.flags.c_contiguous):
raise TypeError(f"{type(self).__name__}: params[{k}] must be a C-contiguous float32 "
"NumPy array or a SequenceDeviceTensor (it is updated in place)")
self.lr = float(lr)
self.options = dict(options)
self.numeric_mode = numeric_mode
self._kind, self._flags, self._fp = optimizer_arguments(self._NAME, self.options)
self.n_total = int(sum(p.size for p in self.params)) # glue: sizes per parameter tensor, not elements
if self.n_total < 1:
raise ValueError(f"{type(self).__name__}: no parameters")
self._res_init([np.zeros(self.n_total, dtype=np.float32) for _ in range(3)]) # glue: the three optimizer state slots
init = self._fp[5]
if init:
self._state[1][:] = np.float32(init)
self.t = 0
# the host scalars' running state (beta1^t, beta2^t, NAdam's mu
# product) after step _sc_t, advanced by the binding each step, so a
# step is O(1) instead of a replay of steps 1 .. t - 1
self._sc = np.ones(3, dtype=np.float32)
self._sc_t = 0
def _pack(self, arrays, what):
if isinstance(arrays, np.ndarray):
arrays = [arrays]
if len(arrays) != len(self.params):
raise ValueError(f"{type(self).__name__}: {len(arrays)} {what}, {len(self.params)} params")
for a, p in zip(arrays, self.params): # glue: per parameter tensor, not elements
if np.shape(a) != p.shape:
raise ValueError(f"{type(self).__name__}: a {what} of shape {np.shape(a)} for a param {p.shape}")
if np.asarray(a).dtype == np.float64:
raise TypeError(f"{type(self).__name__}: float64 {what} are refused")
return np.ascontiguousarray(np.concatenate([np.asarray(a, dtype=np.float32).ravel() for a in arrays])) # glue: packs tensors into the flat buffer, a copy
def _flat_view(self, a):
"""`a` itself, flat, when it already IS the packed buffer (one
C-contiguous float32 array of the param's shape): no copy in or out
(lane sequence-cpu)."""
if len(self.params) != 1:
return None
if isinstance(a, (list, tuple)):
if len(a) != 1:
return None
a = a[0]
if (isinstance(a, np.ndarray) and a.dtype == np.float32 and a.flags.c_contiguous
and a.shape == self.params[0].shape):
return a.reshape(-1)
return None
def _step_resident(self, b, grads):
"""The step with the state on the device (`_ResidentState`)."""
self._check_params()
gs = self._grad_list(grads)
fp = [self.lr] + [float(v) for v in self._fp[:5]] # glue: the five optimizer float arguments
if self._res is None:
self._res_opened(b, b.optimizer_resident_open([self.n_total, self._kind, self._flags], fp),
init_zero=not self._fp[5])
self._res_upload()
# empty tensors hold nothing to move or update
keep = [k for k, p in enumerate(self.params) if p.size > 0] # glue: per parameter tensor, not elements
ps = [self.params[k] for k in keep] # glue: per parameter tensor, not elements
b.optimizer_resident_step(self._res, [p.ctypes.data for p in ps] + [gs[k].ctypes.data for k in keep] # glue: addresses per parameter tensor, not elements
+ [self._sc.ctypes.data],
[len(ps), self._kind, self._flags, self.t, self._sc_t]
+ [int(p.size) for p in ps], fp) # glue: sizes per parameter tensor, not elements
del gs
def _step_device(self, b, grads):
"""The step on device params (`optimizer_resident_step_dev`)."""
if sum(p.size for p in self.params) != self.n_total: # glue: sizes per parameter tensor, not elements
raise ValueError(f"{type(self).__name__}: the params were resized under the optimizer")
gs = _device_grads(b, grads, self.params, type(self).__name__)
fp = [self.lr] + [float(v) for v in self._fp[:5]] # glue: the five optimizer float arguments
if self._res is None:
self._res_opened(b, b.optimizer_resident_open([self.n_total, self._kind, self._flags], fp),
init_zero=not self._fp[5])
self._res_upload()
keep = [k for k, p in enumerate(self.params) if p.size > 0] # glue: per parameter tensor, not elements
b.optimizer_resident_step_dev(self._res, [self.params[k].h for k in keep] + [gs[k].h for k in keep] # glue: handles per parameter tensor, not elements
+ [self._sc.ctypes.data],
[len(keep), self._kind, self._flags, self.t, self._sc_t]
+ [int(self.params[k].size) for k in keep], fp) # glue: sizes per parameter tensor, not elements
del gs
def step(self, grads):
"""One update of every param from `grads` (same shapes, same order)."""
if self._dev_b is not None:
self.t += 1
if getattr(self, "lr_schedule", None) is not None:
self.lr = float(self.lr_schedule.lr_at(self.t))
if self._sc_t >= self.t:
self._sc[:] = 1.0
self._sc_t = 0
try:
self._step_device(self._dev_b, grads)
except BaseException:
self.t -= 1
raise
self._sc_t = self.t
return self
b = _backend.binding("_mojolearn_x_sequence", self.numeric_mode)
if self._res is not None or _resident_binding(b, _SEQ_RESIDENT):
self.t += 1
if getattr(self, "lr_schedule", None) is not None:
self.lr = float(self.lr_schedule.lr_at(self.t))
if self._sc_t >= self.t:
self._sc[:] = 1.0
self._sc_t = 0
try:
self._step_resident(b, grads)
except BaseException:
self.t -= 1
raise
self._sc_t = self.t
return self
flat = self._flat_view(self.params)
g = self._flat_view(grads)
if flat is None or g is None or not flat.flags.writeable or np.shares_memory(flat, g):
flat = None
g = self._pack(grads, "grads")
in_place = flat is not None
if not in_place:
flat = self._pack(self.params, "params")
self.t += 1
if getattr(self, "lr_schedule", None) is not None:
self.lr = float(self.lr_schedule.lr_at(self.t))
if self._sc_t >= self.t:
self._sc[:] = 1.0
self._sc_t = 0
b.optimizer_step([flat.ctypes.data, g.ctypes.data] + [s.ctypes.data for s in self._state] # glue: buffer addresses of the state slots
+ [self._sc.ctypes.data],
[self.n_total, self._kind, self._flags, self.t, self._sc_t],
[self.lr] + [float(v) for v in self._fp[:5]]) # glue: the five optimizer float arguments
self._sc_t = self.t
if not in_place:
off = 0
for p in self.params: # glue: unpacks the flat buffer per parameter tensor, not elements
p.ravel()[:] = flat[off:off + p.size]
off += p.size
return self
def state_dict(self):
return dict(t=self.t, lr=self.lr, state=[s.copy() for s in self.state], options=dict(getattr(self, "options", {})), # glue: copies of the state slots
scalars=(self._sc_t, self._sc.copy()))
def load_state_dict(self, sd):
self.t = int(sd["t"])
self.lr = float(sd["lr"])
for dst, src in zip(self.state, sd["state"]): # glue: the optimizer state slots
dst[:] = np.asarray(src, dtype=np.float32)
# (reading `state` above marked the host copies for upload before
# the next resident step)
# without the scalars (an older state dict) the next step replays
# them once from step 1: the same bits
self._sc[:] = 1.0
self._sc_t = 0
sc = sd.get("scalars")
if sc is not None and 0 <= int(sc[0]) <= self.t:
self._sc_t = int(sc[0])
self._sc[:] = np.asarray(sc[1], dtype=np.float32)
return self
[docs]
class RMSprop(_SeqOptimizer):
"""RMSprop optimizer, as `torch.optim.RMSprop`.
Keeps a running average of squared gradients, ``v = alpha * v + (1 - alpha) * g**2``,
and divides the gradient by ``sqrt(v) + eps``. With ``centered=True`` the squared
running mean of the gradient is subtracted from `v`; with ``momentum > 0`` a
momentum buffer accumulates the scaled gradient. Defaults are torch's.
Parameters
----------
params : ndarray or list of ndarray
Parameters to update in place.
lr : float, default 1e-2
Learning rate.
alpha : float, default 0.99
Smoothing constant of the squared-gradient average.
eps : float, default 1e-8
Term added to the denominator.
weight_decay : float, default 0.0
L2 penalty added to the gradient.
momentum : float, default 0.0
Momentum factor.
centered : bool, default False
Normalize by the estimated gradient variance.
numeric_mode : {'identical', 'deterministic', 'fast'} or None, default None
Numeric mode of the update. None uses the library default.
Attributes
----------
t : int
Number of steps taken.
square_avg : ndarray
Running average of squared gradients, flat over all parameters.
Notes
-----
`params` must be C-contiguous float32 NumPy arrays, updated in place; anything
else raises `TypeError`. `step(grads)`
takes one gradient per parameter, in the same order and shapes; float64 gradients
raise `TypeError`. `lr` is read at every step, so you can change ``opt.lr`` between
steps, or set ``opt.lr_schedule`` to a schedule (`StepLR`, `ExponentialLR`,
`OneCycleLR`, or any object with a one-based `lr_at(t)`), which sets `lr` before
each step. `state_dict()` and `load_state_dict()` save and restore the step count,
learning rate and state buffers.
Not supported: parameter groups, `maximize`, `foreach`/`fused`/`capturable`/
`differentiable`, sparse gradients and float64 parameters.
Examples
--------
>>> import numpy as np, mojolearn
>>> w = np.zeros((64, 8), dtype=np.float32)
>>> opt = mojolearn.RMSprop([w], lr=1e-2)
>>> opt = opt.step([np.ones_like(w)])
"""
_NAME = "rmsprop"
[docs]
def __init__(self, params, lr=1e-2, alpha=0.99, eps=1e-8, weight_decay=0.0, momentum=0.0,
centered=False, numeric_mode=None):
super().__init__(params, lr, dict(alpha=alpha, eps=eps, weight_decay=weight_decay,
momentum=momentum, centered=centered), numeric_mode)
@property
def square_avg(self):
"""Running average of squared gradients, flat over all parameters."""
return self.state[1]
[docs]
class Adagrad(_SeqOptimizer):
"""Adagrad optimizer, as `torch.optim.Adagrad`.
Accumulates squared gradients, ``sum += g**2``, and updates
``p -= clr * g / (sqrt(sum) + eps)`` with ``clr = lr / (1 + (t - 1) * lr_decay)``.
Defaults are torch's.
Parameters
----------
params : ndarray or list of ndarray
Parameters to update in place.
lr : float, default 1e-2
Learning rate.
lr_decay : float, default 0.0
Learning-rate decay per step.
weight_decay : float, default 0.0
L2 penalty added to the gradient.
initial_accumulator_value : float, default 0.0
Starting value of the squared-gradient sum.
eps : float, default 1e-10
Term added to the denominator.
numeric_mode : {'identical', 'deterministic', 'fast'} or None, default None
Numeric mode of the update. None uses the library default.
Attributes
----------
t : int
Number of steps taken.
sum : ndarray
Accumulated squared gradients, flat over all parameters.
Notes
-----
`params` must be C-contiguous float32 NumPy arrays, updated in place; anything
else raises `TypeError`. `step(grads)`
takes one gradient per parameter, in the same order and shapes; float64 gradients
raise `TypeError`. `lr` is read at every step, so you can change ``opt.lr`` between
steps, or set ``opt.lr_schedule`` to a schedule (`StepLR`, `ExponentialLR`,
`OneCycleLR`, or any object with a one-based `lr_at(t)`), which sets `lr` before
each step. `state_dict()` and `load_state_dict()` save and restore the step count,
learning rate and state buffers.
Not supported: parameter groups, `maximize`, `foreach`/`fused`/`capturable`/
`differentiable`, sparse gradients and float64 parameters.
Examples
--------
>>> import numpy as np, mojolearn
>>> w = np.zeros(128, dtype=np.float32)
>>> opt = mojolearn.Adagrad(w, lr=0.1).step(np.ones_like(w))
"""
_NAME = "adagrad"
[docs]
def __init__(self, params, lr=1e-2, lr_decay=0.0, weight_decay=0.0, initial_accumulator_value=0.0,
eps=1e-10, numeric_mode=None):
super().__init__(params, lr, dict(lr_decay=lr_decay, weight_decay=weight_decay,
initial_accumulator_value=initial_accumulator_value, eps=eps),
numeric_mode)
@property
def sum(self):
"""Accumulated squared gradients, flat over all parameters."""
return self.state[1]
[docs]
class Lion(_SeqOptimizer):
"""Lion optimizer (Chen et al., 2023, "Symbolic Discovery of Optimization Algorithms").
Follows the lion-pytorch statement of the update: ``p *= 1 - lr * weight_decay``;
``p -= lr * sign(beta1 * m + (1 - beta1) * g)``; ``m = beta2 * m + (1 - beta2) * g``.
Defaults are the paper's.
Parameters
----------
params : ndarray or list of ndarray
Parameters to update in place.
lr : float, default 1e-4
Learning rate.
betas : tuple of float, default (0.9, 0.99)
Interpolation factor of the update direction and decay of the momentum.
weight_decay : float, default 0.0
Decoupled weight decay.
numeric_mode : {'identical', 'deterministic', 'fast'} or None, default None
Numeric mode of the update. None uses the library default.
Attributes
----------
t : int
Number of steps taken.
exp_avg : ndarray
Momentum buffer, flat over all parameters.
Notes
-----
`params` must be C-contiguous float32 NumPy arrays, updated in place; anything
else raises `TypeError`. `step(grads)`
takes one gradient per parameter, in the same order and shapes; float64 gradients
raise `TypeError`. `lr` is read at every step, so you can change ``opt.lr`` between
steps, or set ``opt.lr_schedule`` to a schedule (`StepLR`, `ExponentialLR`,
`OneCycleLR`, or any object with a one-based `lr_at(t)`), which sets `lr` before
each step. `state_dict()` and `load_state_dict()` save and restore the step count,
learning rate and state buffers.
Not supported: parameter groups, `maximize`, `foreach`/`fused`/`capturable`/
`differentiable`, sparse gradients and float64 parameters.
Examples
--------
>>> import numpy as np, mojolearn
>>> w = np.zeros(128, dtype=np.float32)
>>> opt = mojolearn.Lion([w], lr=1e-4).step([np.ones_like(w)])
"""
_NAME = "lion"
[docs]
def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0, numeric_mode=None):
super().__init__(params, lr, dict(betas=betas, weight_decay=weight_decay), numeric_mode)
@property
def exp_avg(self):
"""Momentum buffer, flat over all parameters."""
return self.state[0]
[docs]
class Adamax(_SeqOptimizer):
"""Adamax optimizer, as `torch.optim.Adamax`.
The infinity-norm variant of Adam: ``m = lerp(m, g, 1 - beta1)``,
``u = max(beta2 * u, |g| + eps)``, ``p -= lr / (1 - beta1**t) * m / u``. Weight
decay is coupled (added to the gradient). Defaults are torch's.
Parameters
----------
params : ndarray or list of ndarray
Parameters to update in place.
lr : float, default 2e-3
Learning rate.
betas : tuple of float, default (0.9, 0.999)
Decay rates of the first moment and the infinity norm, each in [0, 1).
eps : float, default 1e-8
Term added to the norm.
weight_decay : float, default 0.0
L2 penalty added to the gradient.
numeric_mode : {'identical', 'deterministic', 'fast'} or None, default None
Numeric mode of the update. None uses the library default.
Attributes
----------
t : int
Number of steps taken.
Notes
-----
`params` must be C-contiguous float32 NumPy arrays, updated in place; anything
else raises `TypeError`. `step(grads)`
takes one gradient per parameter, in the same order and shapes; float64 gradients
raise `TypeError`. `lr` is read at every step, so you can change ``opt.lr`` between
steps, or set ``opt.lr_schedule`` to a schedule (`StepLR`, `ExponentialLR`,
`OneCycleLR`, or any object with a one-based `lr_at(t)`), which sets `lr` before
each step. `state_dict()` and `load_state_dict()` save and restore the step count,
learning rate and state buffers.
Not supported: parameter groups, `maximize`, `foreach`/`fused`/`capturable`/
`differentiable`, sparse gradients and float64 parameters.
Examples
--------
>>> import numpy as np, mojolearn
>>> w = np.zeros(128, dtype=np.float32)
>>> opt = mojolearn.Adamax([w]).step([np.ones_like(w)])
"""
_NAME = "adamax"
[docs]
def __init__(self, params, lr=2e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0, numeric_mode=None):
super().__init__(params, lr, dict(betas=betas, eps=eps, weight_decay=weight_decay), numeric_mode)
[docs]
class NAdam(_SeqOptimizer):
"""NAdam optimizer, as `torch.optim.NAdam`.
Adam with Nesterov momentum and torch's momentum schedule
``mu_t = beta1 * (1 - 0.5 * 0.96**(t * momentum_decay))``; the denominator is
``sqrt(v / (1 - beta2**t)) + eps``. Weight decay is coupled by default, decoupled
(AdamW style) with ``decoupled_weight_decay=True``. Defaults are torch's.
Parameters
----------
params : ndarray or list of ndarray
Parameters to update in place.
lr : float, default 2e-3
Learning rate.
betas : tuple of float, default (0.9, 0.999)
Decay rates of the first and second moments.
eps : float, default 1e-8
Term added to the denominator.
weight_decay : float, default 0.0
Weight decay.
momentum_decay : float, default 4e-3
Momentum schedule decay. Must be >= 0, otherwise `ValueError`.
decoupled_weight_decay : bool, default False
Apply weight decay to the parameters directly instead of the gradient.
numeric_mode : {'identical', 'deterministic', 'fast'} or None, default None
Numeric mode of the update. None uses the library default.
Attributes
----------
t : int
Number of steps taken.
Notes
-----
`params` must be C-contiguous float32 NumPy arrays, updated in place; anything
else raises `TypeError`. `step(grads)`
takes one gradient per parameter, in the same order and shapes; float64 gradients
raise `TypeError`. `lr` is read at every step, so you can change ``opt.lr`` between
steps, or set ``opt.lr_schedule`` to a schedule (`StepLR`, `ExponentialLR`,
`OneCycleLR`, or any object with a one-based `lr_at(t)`), which sets `lr` before
each step. `state_dict()` and `load_state_dict()` save and restore the step count,
learning rate and state buffers.
Not supported: parameter groups, `maximize`, `foreach`/`fused`/`capturable`/
`differentiable`, sparse gradients and float64 parameters.
Examples
--------
>>> import numpy as np, mojolearn
>>> w = np.zeros(128, dtype=np.float32)
>>> opt = mojolearn.NAdam([w], lr=2e-3).step([np.ones_like(w)])
"""
_NAME = "nadam"
[docs]
def __init__(self, params, lr=2e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0, momentum_decay=4e-3,
decoupled_weight_decay=False, numeric_mode=None):
super().__init__(params, lr, dict(betas=betas, eps=eps, weight_decay=weight_decay,
momentum_decay=momentum_decay,
decoupled_weight_decay=decoupled_weight_decay), numeric_mode)
class Adafactor:
"""`torch.optim.Adafactor` (PyTorch 2.5): relative step size, decoupled
weight decay, a factored second moment (row and column means of g^2) for
matrices and a full one for vectors, the update clipped by its RMS over
d. Parameters are float32 NumPy arrays of 1 or 2 dimensions, updated in
place; defaults are torch's. `sequence/adafactor.mojo`."""
def __init__(self, params, lr=1e-2, beta2_decay=-0.8, eps=(None, 1e-3), d=1.0, weight_decay=0.0,
maximize=False, numeric_mode=None):
if maximize:
raise NotImplementedError("Adafactor: maximize is not implemented")
if isinstance(params, np.ndarray):
params = [params]
self.params = list(params)
self._dev_b = _device_params(self.params, "Adafactor", _AF_DEV)
for k, p in enumerate(self.params): # glue: per parameter tensor, not elements
if self._dev_b is None and (not isinstance(p, np.ndarray) or p.dtype != np.float32
or not p.flags.c_contiguous):
raise TypeError(f"Adafactor: params[{k}] must be a C-contiguous float32 NumPy array "
"or a SequenceDeviceTensor")
if p.ndim not in (1, 2):
raise NotImplementedError("Adafactor: tensors of more than 2 dimensions are not implemented")
if beta2_decay > 0:
raise ValueError("Adafactor: beta2_decay must be <= 0")
if d < 1.0:
raise ValueError(f"Adafactor: clipping threshold d must be >= 1, got {d}")
if lr < 0 or weight_decay < 0:
raise ValueError("Adafactor: lr and weight_decay must be >= 0")
self.lr, self.beta2_decay, self.d, self.weight_decay = float(lr), float(beta2_decay), float(d), float(weight_decay)
eps1, eps2 = eps
self.eps = (float(np.finfo(np.float32).eps) if eps1 is None else float(eps1), float(eps2))
self.numeric_mode = numeric_mode
self._af_state = []
for p in self.params: # glue: per parameter tensor, not elements
if p.ndim == 2:
self._af_state.append(dict(row_var=np.zeros(p.shape[0], np.float32),
col_var=np.zeros(p.shape[1], np.float32)))
else:
self._af_state.append(dict(variance=np.zeros(p.shape[0], np.float32)))
self.t = 0
# lane apple-fast-gap-optim: the second moment on the device
# (`adafactor_resident_*`, a FAST + Apple binding, the default since
# rab10-afresident; -D MOJOLEARN_AF_FAST_RESIDENT_OFF), mirrored on the host on access as
# `_ResidentState` does: `state` downloads once per stale read and the
# next step uploads the host copies a reader may have written
self._af_res = None
self._af_b = None
self._af_fresh = True
self._af_owned = True
self._af_pristine = True
@staticmethod
def _af_slots(st):
return [st["row_var"], st["col_var"]] if "row_var" in st else [st["variance"]]
@property
def state(self):
if self._af_res is not None and not self._af_fresh:
for h, st in zip(self._af_res, self._af_state): # glue: per parameter tensor, not elements
for k, s in enumerate(self._af_slots(st)): # glue: per moment slot (one or two)
self._af_b.optimizer_resident_move(h, k, s.ctypes.data, 0)
self._af_fresh = True
self._af_owned = True
self._af_pristine = False
return self._af_state
@state.setter
def state(self, value):
self._af_state = value
self._af_fresh = True
self._af_owned = True
self._af_pristine = False
def __del__(self):
res, b = getattr(self, "_af_res", None), getattr(self, "_af_b", None)
if res is not None and b is not None:
for h in res: # glue: closes each tensor handle
try:
b.optimizer_resident_close(h)
except Exception:
pass
def _step_resident(self, b, grads):
if self._af_res is None:
res, zeroed = [], True
for p in self.params: # glue: per parameter tensor, not elements
R, C = (p.shape[0], p.shape[1]) if p.ndim == 2 else (p.shape[0], 0)
ret = b.adafactor_resident_open([R, C])
res.append(int(ret[0]))
zeroed = zeroed and len(ret) > 2 and int(ret[2]) == 1
self._af_res, self._af_b = res, b
if zeroed and self._af_pristine:
self._af_owned = False
if self._af_owned:
for h, st in zip(self._af_res, self._af_state): # glue: per parameter tensor, not elements
for k, s in enumerate(self._af_slots(st)): # glue: per moment slot (one or two)
b.optimizer_resident_move(h, k, s.ctypes.data, 1)
self._af_owned = False
self._af_fresh = False
fp = [self.lr, self.beta2_decay, self.eps[0], self.eps[1], self.d, self.weight_decay]
for h, p, g in zip(self._af_res, self.params, grads): # glue: per parameter tensor, not elements
R, C = (p.shape[0], p.shape[1]) if p.ndim == 2 else (p.shape[0], 0)
if self._dev_b is not None:
b.adafactor_resident_step_dev(h, [p.h, g.h], [R, C, self.t], fp)
else:
b.adafactor_resident_step(h, [p.ctypes.data, g.ctypes.data], [R, C, self.t], fp)
def step(self, grads):
if isinstance(grads, np.ndarray):
grads = [grads]
if len(grads) != len(self.params):
raise ValueError(f"Adafactor: {len(grads)} grads, {len(self.params)} params")
self.t += 1
if getattr(self, "lr_schedule", None) is not None:
self.lr = float(self.lr_schedule.lr_at(self.t))
if self._dev_b is not None:
# device params (lane gap-neural-io): updated where they live
try:
self._step_resident(self._dev_b, _device_grads(self._dev_b, grads, self.params, "Adafactor"))
except BaseException:
self.t -= 1
raise
return self
b = _backend.binding("_mojolearn_x_sequence", self.numeric_mode)
if self._af_res is not None or _resident_binding(b, _AF_RESIDENT):
gs = []
for p, g in zip(self.params, grads): # glue: per parameter tensor, not elements
g = np.asarray(g)
if g.shape != p.shape or g.dtype == np.float64:
self.t -= 1
raise ValueError("Adafactor: a grad must be float32 of its param's shape")
gs.append(np.ascontiguousarray(g, dtype=np.float32))
try:
self._step_resident(b, gs)
except BaseException:
self.t -= 1
raise
return self
for p, g, st in zip(self.params, grads, self._af_state): # glue: per parameter tensor, not elements
g = np.asarray(g)
if g.shape != p.shape or g.dtype == np.float64:
raise ValueError("Adafactor: a grad must be float32 of its param's shape")
g = np.ascontiguousarray(g, dtype=np.float32)
if p.ndim == 2:
s1, s2 = st["row_var"], st["col_var"]
R, C = p.shape
else:
s1 = s2 = st["variance"]
R, C = p.shape[0], 0
b.adafactor_step([p.ctypes.data, g.ctypes.data, s1.ctypes.data, s2.ctypes.data], [R, C, self.t],
[self.lr, self.beta2_decay, self.eps[0], self.eps[1], self.d, self.weight_decay])
return self
class LAMB(_SeqOptimizer):
"""LAMB (You et al. 2019, "Large Batch Optimization for Deep Learning"),
timm's `Lamb` statement: optional global gradient-norm clip
(max_grad_norm), Adam moments with bias correction, update
m_hat / (sqrt(v_hat) + eps) + weight_decay p, scaled per tensor by the
trust ratio ||p|| / ||update|| (1 when either is 0; at most 1 with
trust_clip) when weight_decay != 0 or always_adapt. Defaults are timm's."""
def __init__(self, params, lr=1e-3, bias_correction=True, betas=(0.9, 0.999), eps=1e-6,
weight_decay=0.01, grad_averaging=True, max_grad_norm=1.0, trust_clip=False,
always_adapt=False, numeric_mode=None):
if isinstance(params, np.ndarray):
params = [params]
self.params = list(params)
self._dev_b = _device_params(self.params, "LAMB", _LAMB_DEV)
for k, p in enumerate(self.params): # glue: per parameter tensor, not elements
if p.size == 0 or (self._dev_b is None and (not isinstance(p, np.ndarray) or p.dtype != np.float32
or not p.flags.c_contiguous)):
raise TypeError(f"LAMB: params[{k}] must be a non-empty C-contiguous float32 NumPy array "
"or SequenceDeviceTensor")
b1, b2 = betas
if not (0.0 <= b1 < 1.0 and 0.0 <= b2 < 1.0):
raise ValueError("LAMB: betas must lie in [0, 1)")
self.lr, self.betas, self.eps, self.weight_decay = float(lr), (float(b1), float(b2)), float(eps), float(weight_decay)
self.max_grad_norm = None if max_grad_norm is None else float(max_grad_norm)
self.flags = (int(bool(trust_clip)) | 2 * int(bool(always_adapt)) | 4 * int(bool(grad_averaging))
| 8 * int(bool(bias_correction)) | 16 * int(self.max_grad_norm is not None))
self.numeric_mode = numeric_mode
self.n_total = int(sum(p.size for p in self.params)) # glue: sizes per parameter tensor, not elements
self._res_init([np.zeros(self.n_total, dtype=np.float32) for _ in range(2)]) # glue: the two optimizer state slots
self.t = 0
self._sc = np.ones(2, dtype=np.float32) # beta1^_sc_t, beta2^_sc_t
self._sc_t = 0
def step(self, grads):
dev = self._dev_b is not None
b = self._dev_b if dev else _backend.binding("_mojolearn_x_sequence", self.numeric_mode)
resident = dev or self._res is not None or _resident_binding(b, _LAMB_RESIDENT)
if dev:
# device params (lane gap-neural-io): updated where they live
gs = _device_grads(b, grads, self.params, "LAMB")
elif resident:
self._check_params()
gs = self._grad_list(grads)
else:
g = self._pack(grads, "grads")
flat = self._pack(self.params, "params")
self.t += 1
if getattr(self, "lr_schedule", None) is not None:
self.lr = float(self.lr_schedule.lr_at(self.t))
offs = [0]
for p in self.params: # glue: offsets per parameter tensor, not elements
offs.append(offs[-1] + p.size)
if self._sc_t >= self.t:
self._sc[:] = 1.0
self._sc_t = 0
fp = [self.lr, self.betas[0], self.betas[1], self.eps, self.weight_decay,
self.max_grad_norm if self.max_grad_norm is not None else 1.0, float(self._sc_t)]
if resident:
# the moments and the tensor table on the device (`_ResidentState`,
# `sequence/opt_resident.mojo::lamb_resident_step`): each param and
# grad straight into its place, the params straight back
try:
if self._res is None:
self._res_opened(b, b.lamb_resident_open([len(self.params)] + offs))
self._res_upload()
if dev:
b.lamb_resident_step_dev(self._res, [p.h for p in self.params] + [x.h for x in gs] # glue: handles per parameter tensor, not elements
+ [self._sc.ctypes.data], [len(self.params), self.t, self.flags], fp)
else:
b.lamb_resident_step(self._res, [p.ctypes.data for p in self.params] + [x.ctypes.data for x in gs] # glue: addresses per parameter tensor, not elements
+ [self._sc.ctypes.data], [len(self.params), self.t, self.flags], fp)
except BaseException:
self.t -= 1
raise
del gs
else:
b.lamb_step(
[flat.ctypes.data, g.ctypes.data, self._state[0].ctypes.data, self._state[1].ctypes.data,
self._sc.ctypes.data],
[len(self.params), self.t, self.flags] + offs, fp)
if self.flags & 8:
self._sc_t = self.t
if not resident:
off = 0
for p in self.params: # glue: unpacks the flat buffer per parameter tensor, not elements
p.ravel()[:] = flat[off:off + p.size]
off += p.size
return self
assert set(OPTIMIZERS) >= {"rmsprop", "adagrad"}
