optimizer = optax.chain(
    optax.clip_by_global_norm(max_norm=1.0),
    optax.adam(learning_rate=0.001)
)

        __
        __