Looks like Jax fixed their interface: https://github.com/jax-ml/jax/issues/20620
Looks like Jax fixed their interface: jax-ml/jax#20620