[flax]fix jax array type check
# What does this PR do?

Fixes #12584, #12578

On colab the `ModelOutput` class is returning empty tuples for jax arrays. This is because on colab TPU the type of jax array is `jax.interpreters.xla._DeviceArray` and the `is_tensor` function here

https://github.com/huggingface/transformers/blob/2dd9440d0835782e41ae415a68e71fd15051c428/src/transformers/file_utils.py#L1796-L1798

expects `jaxlib.xla_extension.DeviceArray` or `jax.core.Tracer`. If the first argument is an array, the `is_tensor` returns `None` in which case the  `ModelOutput` class expects the first argument to be a key-value container which is not the case here. So at the end, everything becomes `None` and the `ModelOutput` returns an empty tuple.

Instead the `jnp.ndarray` type check works for  jax array types `jaxlib.xla_extension.DeviceArray`, `jax.interpreters.xla._DeviceArray`  and also the `ShardedDeviceArray`

