optimizer = optax.chain(
    optax.clip(max_delta=1.0),  # Clip each element to [-1, 1]
    optax.adam(learning_rate=0.001)
)

        __
        __