pytensor_ml.optim.chain

pytensor_ml.optim.chain(*transforms)

Compose transforms left to right, each reading what the one before it produced.

Every argument has the same type, so a clip composes ahead of a rule as readily as behind it, and the two mean different things. Ahead of the rule the clip sees gradients, so a spike is bounded before it reaches the moment estimates; behind it the clip sees the step the rule already decided on, which an adaptive rule has normalized to roughly its learning rate whatever the gradient was.

stop_the_spike = chain(clip_by_global_norm(1.0), adam(1e-3))
bound_the_move = chain(adam(1e-3), clip_by_global_norm(1.0))

A chain is itself a transform, so one composes into another and the result is flat.

The composed callable owns one set of optimizer-state buffers however many times it is invoked, so two training functions compiled from one chain share its momentum rather than each allocating their own.

Parameters:

*transformsTransform

Applied in order. The first reads whatever the chain is called with – a loss, gradients, or an updates dict – and each one after it reads the previous one’s output.

Returns:

chainedTransform

A transform applying every argument in sequence.

Examples

Clip the gradients before the rule sees them, which is what bounds an exploding gradient rather than the step it produced:

import numpy as np

from pytensor_ml.layers import Input, Linear
from pytensor_ml.loss import SquaredError, supervised_loss
from pytensor_ml.optim import adam, chain, clip_by_global_norm, compile_train

X = Input("X", shape=(None, 4))
loss, target = supervised_loss(Linear("fc", n_in=4, n_out=1)(X), SquaredError())

step = compile_train(loss, chain(clip_by_global_norm(1.0), adam(1e-3)))
loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))

Put a transform after the rule to act on the step instead, which is where a rate or a decay belongs:

from pytensor_ml.optim import adam, chain, clip_by_global_norm, scale

rule = chain(clip_by_global_norm(1.0), adam(1.0), scale(1e-3))