optax.clip

optax.clip(max_delta: jax.typing.ArrayLike) → optax.GradientTransformation

Clips updates element-wise, to be in [-max_delta, +max_delta].

Parameters:

max_delta – The maximum absolute value for each element in the update.

Returns:

A optax.GradientTransformation object.