Remove users of jax.api.* symbols, in preparation for removing the deprecated jax.api name.
Remove users of jax.api.* symbols, in preparation for removing the deprecated jax.api name.

In most cases, use the public jax.* name instead.

PiperOrigin-RevId: 395804438

google/flax,google/flax