[JAX] Replace uses of x.to_py() on JAX arrays with numpy.asarray(x).
[JAX] Replace uses of x.to_py() on JAX arrays with numpy.asarray(x).

.to_py() was something of an accidental export from the JAX API, and the same thing can be obtained via numpy.asarray(...) or jax.device_get(...).

PiperOrigin-RevId: 470014108

google/paxml,google/paxml