flax

[JAX] Replace uses of `isinstance(x, jax.numpy.ndarray)` with `isinstance(x, (numpy.ndarray, jax.numpy.ndarray))` where they were leading to test failures.

#1549

Merged