Add Grain integration. (#21494)

* Add Grain integration.

* Fix CI and resolve comments from gemini-code-assist.

* Add ragged attrs to `KerasTensor` and update `data_adapter_utils.py`.

* Simplify the impl for getting `batch_size` and `output_signature`.

* Add `builtin_prefetch` to `TFDatasetAdapter` and `TorchDataLoaderAdapter`.

* Use `.map` instead of the internal grain functions.

Author: james77777778

keras/src/backend/common/keras_tensor.py

@@ -35,9 +35,17 @@ def __init__(
         ragged=False,
         record_history=True,
         name=None,
+        **kwargs,
     ):
         from keras.src import backend
 
+        ragged_rank = kwargs.pop("ragged_rank", None)
+        row_splits_dtype = kwargs.pop("row_splits_dtype", None)
+        if kwargs:
+            raise TypeError(
+                f"Unexpected keyword arguments: {', '.join(kwargs.keys())}"
+            )
+
         self._shape = backend.standardize_shape(shape)
         self._dtype = backend.standardize_dtype(dtype)
         self._sparse = bool(sparse)
@@ -47,6 +55,14 @@ def __init__(
                 "KerasTensor cannot have `sparse=True` and `ragged=True` at "
                 "the same time."
             )
+        self._ragged_rank = (
+            int(ragged_rank) if ragged_rank is not None else None
+        )
+        self._row_splits_dtype = (
+            backend.standardize_dtype(row_splits_dtype)
+            if row_splits_dtype is not None
+            else None
+        )
         self.name = name or auto_name(self.__class__.__name__)
         self.record_history = record_history
 
@@ -83,6 +99,28 @@ def sparse(self, value):
             "create a new instance of KerasTensor for this."
         )
 
+    @property
+    def ragged_rank(self):
+        return self._ragged_rank
+
+    @ragged_rank.setter
+    def ragged_rank(self, value):
+        raise AttributeError(
+            "The `ragged_rank` attribute of KerasTensor is immutable. One "
+            "should create a new instance of KerasTensor for this."
+        )
+
+    @property
+    def row_splits_dtype(self):
+        return self._row_splits_dtype
+
+    @row_splits_dtype.setter
+    def row_splits_dtype(self, value):
+        raise AttributeError(
+            "The `row_splits_dtype` attribute of KerasTensor is immutable. One "
+            "should create a new instance of KerasTensor for this."
+        )
+
     @property
     def ragged(self):
         return self._ragged

keras/src/backend/jax/trainer.py

@@ -1026,10 +1026,12 @@ def _get_iterator(self):
         distribution = distribution_lib.distribution()
         if distribution is not None:
             return self._get_distributed_iterator(distribution)
-
-        return self._prefetch_numpy_iterator(
-            self.data_adapter.get_jax_iterator()
-        )
+        if self.data_adapter.builtin_prefetch:
+            return self.data_adapter.get_jax_iterator()
+        else:
+            return self._prefetch_numpy_iterator(
+                self.data_adapter.get_jax_iterator()
+            )
 
     def _get_distributed_iterator(self, distribution):
         """Lazily compute layouts to reduce host to device transfer latency."""

keras/src/trainers/data_adapters/__init__.py

@@ -8,6 +8,9 @@
 from keras.src.trainers.data_adapters.generator_data_adapter import (
     GeneratorDataAdapter,
 )
+from keras.src.trainers.data_adapters.grain_dataset_adapter import (
+    GrainDatasetAdapter,
+)
 from keras.src.trainers.data_adapters.py_dataset_adapter import PyDatasetAdapter
 from keras.src.trainers.data_adapters.tf_dataset_adapter import TFDatasetAdapter
 from keras.src.trainers.data_adapters.torch_data_loader_adapter import (
@@ -111,6 +114,32 @@ def get_data_adapter(
         #     "data `x` was provided as a torch DataLoader. The DataLoader "
         #     "is expected to already be shuffled."
         # )
+    elif is_grain_dataset(x):
+        if y is not None:
+            raise_unsupported_arg(
+                "y", "the targets", "grain.Dataset and grain.DataLoader"
+            )
+        if sample_weight is not None:
+            raise_unsupported_arg(
+                "sample_weights",
+                "the sample weights",
+                "grain.Dataset and grain.DataLoader",
+            )
+        if class_weight is not None:
+            raise ValueError(
+                "Argument `class_weight` is not supported for grain.Dataset "
+                f"and grain.DataLoader inputs. You can modify your "
+                "`__getitem__ ` method to return input tensor, label and "
+                "class_weight. "
+                f"Received: class_weight={class_weight}"
+            )
+        return GrainDatasetAdapter(x)
+        # TODO: should we warn or not?
+        # warnings.warn(
+        #     "`shuffle=True` was passed, but will be ignored since the "
+        #     "data `x` was provided as a grain dataset. The grain dataset "
+        #     "is expected to already be shuffled."
+        # )
     elif isinstance(x, types.GeneratorType):
         if y is not None:
             raise_unsupported_arg("y", "the targets", "PyDataset")
@@ -162,3 +191,15 @@ def is_torch_dataloader(x):
             ):
                 return True
     return False
+
+
+def is_grain_dataset(x):
+    if hasattr(x, "__class__"):
+        for parent in x.__class__.__mro__:
+            if parent.__name__ in (
+                "MapDataset",
+                "IterDataset",
+                "DataLoader",
+            ) and "grain" in str(parent.__module__):
+                return True
+    return False

keras/src/trainers/data_adapters/data_adapter.py

@@ -46,6 +46,21 @@ def get_torch_dataloader(self):
         """
         raise NotImplementedError
 
+    @property
+    def builtin_prefetch(self):
+        """Whether the DataAdapter has built-in prefetching capabilities.
+
+        Prefetching is an optimization technique where data is loaded and
+        prepared in advance while the model is processing the current batch,
+        reducing training time by overlapping data loading with computation.
+
+        Returns:
+            bool: True if the DataAdapter implements its own prefetching
+            mechanism and handles data loading asynchronously. False if the
+            caller should implement prefetching externally.
+        """
+        return False
+
     @property
     def num_batches(self):
         """Return the size (number of batches) for the dataset created.

keras/src/trainers/data_adapters/data_adapter_utils.py

@@ -133,18 +133,25 @@ def class_weight_to_sample_weights(y, class_weight):
     return sample_weight
 
 
-def get_tensor_spec(batches):
-    """Return the common tensor spec for a list of batches.
+def get_keras_tensor_spec(batches):
+    """Return the KerasTensor spec for a list of batches.
+
+    The spec is represented using `KerasTensor` which could handle dense, sparse
+    or ragged tensors.
 
     Args:
         batches: list of structures of tensors. The structures must be
             identical, but the shape at each leaf may be different.
-    Returns: the common tensor spec for all the batches.
+
+    Returns:
+        A nested structure of `KerasTensor`.
     """
-    from keras.src.utils.module_utils import tensorflow as tf
 
     def get_single_tensor_spec(*tensors):
         x = tensors[0]
+        if not hasattr(x, "shape"):
+            # Try to convert to a numpy array.
+            x = np.array(x)
         rank = len(x.shape)
         if rank < 1:
             raise ValueError(
@@ -164,28 +171,72 @@ def get_single_tensor_spec(*tensors):
         for dims in zip(*[list(x.shape) for x in tensors]):
             dims_set = set(dims)
             shape.append(dims_set.pop() if len(dims_set) == 1 else None)
-        shape[0] = None  # batch size may not be static
 
         dtype = backend.standardize_dtype(x.dtype)
-        if isinstance(x, tf.RaggedTensor):
-            return tf.RaggedTensorSpec(
+        if is_tensorflow_ragged(x):
+            return backend.KerasTensor(
                 shape=shape,
                 dtype=dtype,
+                ragged=True,
                 ragged_rank=x.ragged_rank,
                 row_splits_dtype=x.row_splits.dtype,
             )
-        if (
-            isinstance(x, tf.SparseTensor)
-            or is_scipy_sparse(x)
-            or is_jax_sparse(x)
-        ):
-            return tf.SparseTensorSpec(shape=shape, dtype=dtype)
+        if is_tensorflow_sparse(x) or is_scipy_sparse(x) or is_jax_sparse(x):
+            return backend.KerasTensor(shape=shape, dtype=dtype, sparse=True)
         else:
-            return tf.TensorSpec(shape=shape, dtype=dtype)
+            return backend.KerasTensor(shape=shape, dtype=dtype)
 
     return tree.map_structure(get_single_tensor_spec, *batches)
 
 
+def convert_to_tf_tensor_spec(keras_tensor, batch_axis_to_none=True):
+    """Convert a KerasTensor to a TensorSpec.
+
+    Args:
+        keras_tensor: A KerasTensor instance.
+        batch_axis_to_none: If `True`, the batch axis of the returned
+            tensor spec will be set to None. Defaults to `True`.
+    """
+    from keras.src.utils.module_utils import tensorflow as tf
+
+    if not isinstance(keras_tensor, backend.KerasTensor):
+        raise TypeError(
+            f"Expected a KerasTensor, but got {keras_tensor} of type "
+            f"{type(keras_tensor)}."
+        )
+    shape = list(keras_tensor.shape)
+    if batch_axis_to_none:
+        shape[0] = None
+    if keras_tensor.ragged:
+        return tf.RaggedTensorSpec(
+            shape=shape,
+            dtype=keras_tensor.dtype,
+            ragged_rank=keras_tensor.ragged_rank,
+            row_splits_dtype=keras_tensor.row_splits_dtype,
+        )
+    elif keras_tensor.sparse:
+        return tf.SparseTensorSpec(shape=shape, dtype=keras_tensor.dtype)
+    else:
+        return tf.TensorSpec(shape=shape, dtype=keras_tensor.dtype)
+
+
+def get_tensor_spec(batches):
+    """Return the common tensor spec for a list of batches.
+
+    The spec is represented using `tf.TensorSpec`, `tf.SparseTensorSpec` and
+    `tf.RaggedTensorSpec`.
+
+    Args:
+        batches: list of structures of tensors. The structures must be
+            identical, but the shape at each leaf may be different.
+
+    Returns:
+        A common tensor spec.
+    """
+    tensor_specs = get_keras_tensor_spec(batches)
+    return tree.map_structure(convert_to_tf_tensor_spec, tensor_specs)
+
+
 def get_jax_iterator(iterable):
     import jax
     import jax.experimental.sparse as jax_sparse