[JAX] Update JAX users in preparation for a change that makes iteration over a JAX array return JAX arrays, instead of NumPy arrays.

[copybara-service]

[JAX] Update JAX users in preparation for a change that makes iteration over a JAX array return JAX arrays, instead of NumPy arrays.

See jax-ml/jax#8043 for context as to why we are making this change.

The upshot for most users is that the values returned by iteration over a JAX array are now themselves JAX arrays, with the semantics of JAX arrays, which sometimes differ from the semantics of NumPy scalars and arrays. In particular:

• Unlike NumPy scalars 0-dimensional JAX arrays are not hashable. This change updates users to call .tolist() or np.asarray(...) when the output of iterating over a JAX array is hashed, used as a dictionary key, or passed to set(...). In some instances, we can just call numpy functions instead of jax.numpy functions to build the array in the first place.
• This change confuses Pandas and PIL when a JAX array is converted to a Pandas dataframe or a PIL image. For now, cast JAX arrays to a NumPy array first before passing them into those libraries.
• We now need to use numpy.testing.assert_array_equal instead of numpy.testing.assert_equal to compare JAX arrays.