keras/src/trainers/data_adapters/generator_data_adapter.py/GeneratorDataAdapter/get_jax_iterator
class GeneratorDataAdapter:    def get_jax_iterator(self):
        from keras.src.backend.jax.core import convert_to_tensor

        def convert_to_jax(x):
            if data_adapter_utils.is_scipy_sparse(x):
                return data_adapter_utils.scipy_sparse_to_jax_sparse(x)
            elif data_adapter_utils.is_tensorflow_sparse(x):
                return data_adapter_utils.tf_sparse_to_jax_sparse(x)
            return convert_to_tensor(x)

        for batch in self.generator:
            yield tree.map_structure(convert_to_jax, batch)