tx = optax.chain(
    optax.clip_by_global_norm(0.5),  # Aggressive clipping
    optax.sgd(learning_rate=0.01, momentum=0.9)  # SGD is more stable
)

        __
        __