Relax test tolerance a trivial amount (I think).
Relax test tolerance a trivial amount (I think).

I noticed that this test failed when https://github.com/google/jax/pull/3370 is enabled by default. I think it's just a harmless numerical difference.

PiperOrigin-RevId: 326723573

deepmind/optax,deepmind/optax