Optimization

Training

compile_train(loss, rule, *[, parameters, ...])

Compile a one-step training function from a loss graph and an update rule.

Update rules

sgd([learning_rate, momentum, nesterov, ...])	Stochastic gradient descent, optionally with momentum.
adam([learning_rate, beta1, beta2, epsilon, ...])	Adam optimizer.
adamw([learning_rate, weight_decay, beta1, ...])	AdamW optimizer (Adam with decoupled weight decay).
nadam([learning_rate, beta1, beta2, ...])	Nadam optimizer (Adam with Nesterov momentum).
adamax([learning_rate, beta1, beta2, ...])	AdaMax optimizer (Adam with an infinity-norm denominator).
rmsprop([learning_rate, rho, momentum, ...])	RMSProp optimizer.
rprop([learning_rate, eta_minus, eta_plus, ...])	Rprop optimizer (resilient backpropagation).
adagrad([learning_rate, epsilon, namespace])	AdaGrad optimizer.
adadelta([learning_rate, rho, epsilon, ...])	AdaDelta optimizer.

Transforms

chain(*transforms)	Compose transforms left to right, each reading what the one before it produced.
scale(factor)	Scale each step by a constant factor.
scale_by_schedule(schedule, *[, namespace])	Scale each step by a schedule read off a training clock of its own.
add_weight_decay([weight_decay, mask])	Subtract a weight-decay term \(\lambda p\).
trace([decay, nesterov, namespace])	Accumulate into a velocity buffer (classical or Nesterov momentum).
clip_by_global_norm([max_norm])	Rescale everything by a single factor so its global L2 norm does not exceed max_norm.
clip_by_value([min_value, max_value])	Clamp every value element-wise into [min_value, max_value].

Guards and policies

skip_if(rule[, condition, ...])	Throw away any step condition rejects, leaving the parameters and the optimizer state as they were.
apply_if_finite(rule, *[, ...])	Skip any step that would write a non-finite parameter, leaving the parameters and the optimizer state as they were.
nonfinite()	Throw the step away when any parameter the rule would write is inf or NaN.
large_step(max_norm)	Throw the step away when the global L2 norm of the steps reaches max_norm.
reduce_on_plateau(rule, scale, *[, factor, ...])	Cut scale by factor once the loss has stopped improving.
SkipCondition(decide[, reason])	A rule for throwing a training step away, paired with the phrase naming why.

Schedules

Every step count a schedule takes – total_steps, transition_begin, decay_every, and the boundaries of join_schedules() – accepts a symbolic value as readily as a literal one, so a horizon can be written in terms of the data it will run over. With X the minibatch and n_samples the size of the dataset, total_steps=10 * (n_samples // X.shape[0]) is ten epochs. The same holds for the patience, cooldown, and accumulation_size of reduce_on_plateau() and the max_norm of large_step() and clip_by_global_norm(), and the bounds of clip_by_value(). A count whose value is not known when the graph is built carries its check into the graph instead. That check rides on the value, so it runs wherever the value does: not under jax, which drops assertions, and not where a rewrite has eliminated the value’s last consumer.

constant_schedule(learning_rate)	Hold the learning rate at learning_rate forever.
linear_schedule(learning_rate, total_steps)	Move the learning rate from learning_rate to final_learning_rate at a constant rate.
linear_onecycle_schedule(peak_value, total_steps)	Move the learning rate through three linear phases over one cycle.
cosine_schedule(learning_rate, total_steps)	Move the learning rate from learning_rate to final_learning_rate along a half cosine.
exponential_schedule(learning_rate, ...[, ...])	Move the learning rate from learning_rate to final_learning_rate by a constant factor per step.
polynomial_schedule(learning_rate, total_steps)	Move the learning rate from learning_rate to final_learning_rate along a power curve.
step_decay(learning_rate, *, decay_every[, ...])	Multiply the learning rate by decay_factor every decay_every steps.
join_schedules(schedules, boundaries)	Run schedules one after another, switching at boundaries.

Building blocks

counter(name)	Return the training clock a component counts its own steps on.
get_gradients(loss_or_gradients, parameters)	Return gradients of the loss with respect to parameters, or pass through precomputed gradients.
reuses_state(builds_updates)	Give builds_updates a private set of optimizer-state buffers, reused on every invocation.
scalar_state(name[, fill_value])	Return a floatX scalar shared variable, reused across invocations of a rule wrapped in reuses_state().
state_for(parameter, slot[, fill_value])	Return the optimizer-state shared variable shaped and typed like parameter.
steps_of(updates, parameters)	Return the amount each parameter's entry moves it by.
to_floatx(value)	Return value at the current floatX, casting only a variable stored at something else.
to_updates(loss_gradients_or_updates, parameters)	Return loss_gradients_or_updates as an updates dict, differentiating a loss if that is what it is.

Transform	What every optimizer, clip, and schedule in this module is: a callable taking a loss, gradients, or an updates dict, along with the parameters, and returning the updates dict that moves them.
Updates	Pytensor's native updates contract, and the single currency every transform here speaks.
Gradients	Updates carrying gradients: updates[parameter] - parameter is the gradient \(g\) itself.
Steps	Updates carrying steps: updates[parameter] - parameter is the move a rule decided on.
Schedule	A learning-rate schedule: symbolic step count in, scalar learning rate out.
Rate	A rate a rule multiplies into its step.
LearningRate	What an optimizer alias accepts as its rate, adding a schedule that drives it on-graph.

Decision

Reads the step a rule has proposed and returns a scalar boolean graph: True to throw the step away.

Low-level update functions

sgd_updates(loss_gradients_or_updates, ...)	Vanilla stochastic gradient descent: \(p \leftarrow p - \eta g\).
adam_updates(loss_gradients_or_updates, ...)	Adam optimizer.
adamw_updates(loss_gradients_or_updates, ...)	AdamW: Adam with decoupled weight decay applied directly to the parameter, not the gradient.
nadam_updates(loss_gradients_or_updates, ...)	Nadam: Adam with Nesterov momentum applied to the first-moment estimate.
adamax_updates(loss_gradients_or_updates, ...)	AdaMax: Adam variant using an exponentially weighted infinity norm instead of the second moment.
rmsprop_updates(loss_gradients_or_updates, ...)	RMSProp: per-parameter learning rate scaled by a decaying average of squared gradients.
rprop_updates(loss_gradients_or_updates, ...)	Rprop: resilient backpropagation, stepping by a per-parameter magnitude that adapts to gradient-sign agreement and ignores gradient magnitude.
adagrad_updates(loss_gradients_or_updates, ...)	AdaGrad: per-parameter learning rate scaled by the inverse root of accumulated squared gradients.
adadelta_updates(loss_gradients_or_updates, ...)	AdaDelta: AdaGrad variant with a decaying window of squared gradients and squared updates.