Optim

substrax.optim turns one specification into an optax transformation or an NNX optimizer, and reads back the learning rate the last update applied.

import optax
from flax import nnx

from substrax.optim import (
    EXCLUDE_BIAS_AND_NORM_SCALE,
    OptimizerConfig,
    create_optimizer,
    current_learning_rate,
)

config = OptimizerConfig(
    optimizer_type="adamw",
    learning_rate=optax.warmup_cosine_decay_schedule(0.0, 3e-4, warmup_steps=500, decay_steps=10_000),
    weight_decay=0.05,
    weight_decay_filter=EXCLUDE_BIAS_AND_NORM_SCALE,
    gradient_clip_norm=1.0,
)
optimizer = create_optimizer(model, config)

grads = nnx.grad(loss)(model, batch)
optimizer.update(model, grads)
print(float(current_learning_rate(optimizer)))

Name	Behaviour
OptimizerConfig	Frozen, keyword-only; optimizer_type is one of adam, adamw, sgd, rmsprop, adagrad, lamb, radam, nadam (adam with Nesterov momentum), learning_rate a positive constant or an optax.Schedule, b1, b2, eps, momentum (sgd and rmsprop), weight_decay (adamw and lamb), weight_decay_filter, one of gradient_clip_norm and gradient_clip_value (positive), and wrt; anything else is refused when the configuration is constructed
weight_decay_filter	A Flax NNX filter naming the parameters that are decayed, the polarity of optax's mask; None decays every leaf in wrt, as torch.optim.AdamW(model.parameters()) does; EXCLUDE_BIAS_AND_NORM_SCALE skips biases and normalisation scales, the Hugging Face and timm convention; a string raises, because NNX reads it as a tag, not a path
weight_decay_mask(model, config)	The filter as a static tree of Python booleans over nnx.state(model, config.wrt); a filter that selects nothing raises instead of becoming optax's silent no-op
create_transformation(model, config)	Clipping first, then the base alias wrapped in optax.inject_hyperparams; a schedule is the base optimizer's learning rate, never a scale_by_schedule placed before an adaptive optimizer, whose normalisation would cancel it
create_optimizer(model, config)	nnx.Optimizer(model, create_transformation(model, config), wrt=config.wrt)
current_learning_rate(optimizer)	The rate the last update applied, a scalar array on device; the schedule at step zero before the first update, and unchanged by a step optax.apply_if_finite skips

The optimizer-state tree has the same structure whatever the filter, so a checkpoint written with one filter restores into an optimizer built with another.

substrax.optim

One optimizer specification, one builder over optax, and the learning rate read on device.

OptimizerConfig names the optimizer in optax's terms; create_transformation and create_optimizer build it, passing a schedule as the base optimizer's learning rate and the weight-decay filter as a static mask; current_learning_rate reads the rate the last update applied. EXCLUDE_BIAS_AND_NORM_SCALE is the filter that decays neither biases nor normalisation scales. update_with_line_search steps an optimizer that searches along its direction (optax.lbfgs); switch_at hands over from one transformation to another after a fixed number of updates; with_strong_state_types keeps a transformation's state types fixed across updates, so a jitted step compiles once.

EXCLUDE_BIAS_AND_NORM_SCALE = _is_neither_bias_nor_scale module-attribute

Decay every parameter except biases and normalisation scales, as Hugging Face and timm do.

OptimizerConfig(*, optimizer_type='adam', learning_rate, b1=0.9, b2=0.999, eps=1e-08, momentum=None, weight_decay=0.0, weight_decay_filter=None, gradient_clip_norm=None, gradient_clip_value=None, wrt=nnx.Param) dataclass

What an optimizer is, in optax's terms.

A schedule is passed as the base optimizer's learning rate, never as a scale_by_schedule placed before an adaptive optimizer, whose normalisation would cancel the scaling. The names map onto other libraries as follows.

Field	PyTorch	Flax NNX
learning_rate	lr	the base optax alias's learning_rate
b1, b2	betas	b1, b2
weight_decay	weight_decay on a parameter group	optax.adamw(weight_decay=...)
weight_decay_filter	the parameter group	optax.adamw(mask=...), True where decayed
wrt	model.parameters()	nnx.Optimizer(model, tx, wrt=...)

Attributes:

Name	Type	Description
optimizer_type	OptimizerType	The optax alias; nadam is adam with Nesterov momentum.
learning_rate	float | Schedule	A positive constant, or an optax.Schedule of the step count.
b1	float	First-moment decay of the Adam family and LAMB.
b2	float	Second-moment decay of the Adam family and LAMB.
eps	float	The denominator offset.
momentum	float | None	Momentum of sgd and rmsprop; None for no momentum.
weight_decay	float	Decoupled weight decay, adamw and lamb only.
weight_decay_filter	Filter | None	A Flax NNX filter naming the parameters that are decayed, the polarity of optax's mask; None decays every leaf in wrt, as torch.optim.AdamW(model.parameters()) does. EXCLUDE_BIAS_AND_NORM_SCALE is the Hugging Face and timm convention.
gradient_clip_norm	float | None	Global-norm clipping applied before the update, or None.
gradient_clip_value	float | None	Element-wise clipping applied before the update, or None.
wrt	Filter	The filter naming the parameters the optimizer updates.

__post_init__()

Refuse a specification optax would silently misread.

Raises:

Type	Description
TypeError	If weight_decay_filter is a string, which NNX reads as a tag, not a path.
ValueError	If optimizer_type is not one built here, both clip fields are set or a clip is not positive, weight_decay is negative or set for an optimizer without decoupled decay, momentum is set for an optimizer that has none, or a constant learning_rate is not positive.

create_optimizer(model, config)

Build an nnx.Optimizer over create_transformation for config.wrt.

Parameters:

Name	Type	Description	Default
model	M	The model the optimizer updates.	required
config	OptimizerConfig	The specification.	required

Returns:

Type	Description
Optimizer[M]	The optimizer; step it with optimizer.update(model, grads).

create_transformation(model, config)

Build the optax transformation config describes for model.

Clipping, when configured, comes first; the base alias follows, wrapped in optax.inject_hyperparams so its learning rate, a schedule of optax's own step count when one is given, can be read on device with current_learning_rate. The weight-decay filter becomes a mask of Python booleans over the parameter tree, never optimizer state, so the state has the same structure whatever the filter.

Parameters:

Name	Type	Description	Default
model	Module	The model the optimizer updates, which fixes the mask's tree.	required
config	OptimizerConfig	The specification.	required

Returns:

Type	Description
GradientTransformation	The transformation.

Raises:

Type	Description
ValueError	If config.optimizer_type is not an optax alias this builder knows, or the weight-decay filter selects no parameter.

mask_callable(config)

Return the callable optax evaluates on the parameter tree to get the decay mask.

Parameters:

Name	Type	Description	Default
config	OptimizerConfig	The specification holding weight_decay_filter.	required

Returns:

Type	Description
Callable[[Params], PyTree]	A function from a parameter tree to a tree of Python booleans of the same structure.

mask_from_filter(params, chosen)

Map an NNX filter over a parameter tree, giving a Python boolean per leaf.

The tree may be an nnx.State or the pure tree optax hands a callable mask; a trailing value attribute key of a Variable is not part of the parameter's path.

Parameters:

Name	Type	Description	Default
params	PyTree	The parameter tree.	required
chosen	Filter	The filter, evaluated on each leaf's path and value.	required

Returns:

Type	Description
PyTree	A tree of the same structure with a bool at every leaf.

weight_decay_mask(model, config)

Return the optax mask for config.weight_decay_filter over nnx.state(model, wrt).

The mask is the parameters' pure tree (nested dicts, as nnx.to_pure_dict gives it) with a static Python boolean at every leaf, True where the filter selects the parameter. None selects every leaf. The builder passes optax the callable form, mask_callable, which maps the same filter over whatever parameter tree optax hands it.

Parameters:

Name	Type	Description	Default
model	Module	The model the optimizer updates.	required
config	OptimizerConfig	The specification holding wrt and weight_decay_filter.	required

Returns:

Type	Description
PyTree	The mask tree.

Raises:

Type	Description
ValueError	If the filter selects no parameter, which optax would accept as a silent no-op.

current_learning_rate(optimizer)

Return the learning rate the last update applied, as an array on device.

create_transformation wraps the base alias in optax.inject_hyperparams, which stores the rate it evaluated at optax's own step count. Before the first update the value is the schedule at step zero; a step a wrapper such as optax.apply_if_finite skips leaves it unchanged, because optax's count does not advance.

Parameters:

Name	Type	Description	Default
optimizer	Optimizer[M]	An optimizer built over create_transformation.	required

Returns:

Type	Description
Array	The learning rate, a scalar array.

Raises:

Type	Description
ValueError	If the optimizer's state holds no injected learning rate.

switch_at(first, second, *, step)

The first transformation for step updates, then second.

Each runs under optax.conditionally_transform on the update count, so the switch is a branch of the traced step rather than a Python decision, and each keeps its own state from initialisation. Extra arguments reach a transformation that takes them (optax.lbfgs's value, grad and value_fn).

Parameters:

Name	Type	Description	Default
first	GradientTransformation	The transformation for the first step updates, such as optax.adam.	required
second	GradientTransformation	The transformation from update step on, such as optax.lbfgs().	required
step	int	The number of updates first makes.	required

Returns:

Type	Description
GradientTransformationExtraArgs	The combined transformation, its initial state strongly typed.

Raises:

Type	Description
ValueError	If step is not positive.

update_with_line_search(model, optimizer, loss_fn, *, objective_changed)

Update model with an optimizer that evaluates the loss along its search direction.

The loss, its gradient and the loss as a function of the optimizer.wrt parameters are passed to optimizer.update as value, grad and value_fn, the extra arguments optax.lbfgs and any transformation built on it read. When objective_changed is false and the state holds the value and gradient the previous line search computed at the current parameters, those are used; otherwise they are evaluated. objective_changed is traced, so either answer runs the same compiled step. Under vmap (an ensemble of models) the choice depends on each member's stored value, and a lax.cond over a batched predicate runs both branches, so each step evaluates the loss as it would with objective_changed true: the result is the same and the saving is lost. The loss is taken in the parameters' dtype: optax's line search keeps its values in that dtype, and a wider loss (a float32 model's loss computed in float64 under x64) would make the branches of its lax.cond disagree.

Parameters:

Name	Type	Description	Default
model	M	The model; updated in place.	required
optimizer	Optimizer[M]	An nnx.Optimizer over model whose transformation takes the extra arguments and whose state is strongly typed, such as with_strong_state_types(optax.lbfgs()) or a switch_at ending in it.	required
loss_fn	Callable[[M], Array]	The scalar loss of the model.	required
objective_changed	bool | Array	Whether loss_fn is a different objective from the previous call's (a new batch, resampled points), or the parameters moved outside the line search since it; False only for the same objective at the point it accepted.	required

Returns:

Type	Description
Array	The loss before the update, in the parameters' dtype.

Raises:

Type	Description
ValueError	If the parameters are of more than one floating dtype, which leaves the line search no single dtype to work in, or if the optimizer state has weakly typed leaves, which would compile the step again once its first update types them.

with_strong_state_types(transformation)

transformation with every weakly typed leaf of its initial state strongly typed.

The values and dtypes are unchanged; only the weak type goes, so the state a jitted step receives has the types its update returns. Updates are transformation's own, extra arguments included.

Parameters:

Name	Type	Description	Default
transformation	GradientTransformation	The optax transformation, such as optax.lbfgs().	required

Returns:

Type	Description
GradientTransformationExtraArgs	The transformation with a strongly typed initial state.