Ensure values passed to jax.numpy functions are arrays rather than lists.
Ensure values passed to jax.numpy functions are arrays rather than lists.

Why? This will soon be a requirement in JAX; see https://github.com/google/jax/issues/7737

PiperOrigin-RevId: 394106413

google/trax,google/trax