Avoid deprecated JAX API

Commit

94 days ago

Avoid deprecated JAX API `jax.core.get_opaque_trace_state` is deprecated as of JAX v0.10.0 (https://github.com/jax-ml/jax/pull/36725). PiperOrigin-RevId: 899625682