optax.scale_gradient

optax.scale_gradient(

inputs: optax.ArrayTree,

scale: jax.typing.ArrayLike,

) → optax.ArrayTree

Scales gradients for the backwards pass.

Parameters:

• inputs – A nested array.

• scale – The scale factor for the gradient on the backwards pass.

Returns:

An array of the same structure as inputs, with scaled backward gradient.