keras/src/trainers/data_adapters/data_adapter_utils.py/is_jax_array
def is_jax_array(value):
    if hasattr(value, "__class__"):
        for parent in value.__class__.__mro__:
            if parent.__name__ == "Array" and str(parent.__module__) == "jax":
                return True
    return is_jax_sparse(value)  # JAX sparse arrays do not extend jax.Array