Improve tf.RaggedTensor support in DataAdapters. (#21170)

Previously, only 2D Tensorflow ragged tensors were supported. This adds support for any rank.

Also added tests for ragged tensors with `GeneratorDataAdapter`.

Author: hertschuh

keras/src/trainers/data_adapters/data_adapter_utils.py

@@ -168,7 +168,12 @@ def get_single_tensor_spec(*tensors):
 
         dtype = backend.standardize_dtype(x.dtype)
         if isinstance(x, tf.RaggedTensor):
-            return tf.RaggedTensorSpec(shape=shape, dtype=dtype)
+            return tf.RaggedTensorSpec(
+                shape=shape,
+                dtype=dtype,
+                ragged_rank=x.ragged_rank,
+                row_splits_dtype=x.row_splits.dtype,
+            )
         if (
             isinstance(x, tf.SparseTensor)
             or is_scipy_sparse(x)

keras/src/trainers/data_adapters/generator_data_adapter_test.py

@@ -166,7 +166,7 @@ def generator():
         not backend.SUPPORTS_SPARSE_TENSORS,
         reason="Backend does not support sparse tensors",
     )
-    def test_scipy_sparse_tensors(self, generator_type):
+    def test_sparse_tensors(self, generator_type):
         if generator_type == "tf":
             x = tf.SparseTensor([[0, 0], [1, 2]], [1.0, 2.0], (2, 4))
             y = tf.SparseTensor([[0, 0], [1, 1]], [3.0, 4.0], (2, 2))
@@ -197,3 +197,33 @@ def generate():
             self.assertIsInstance(by, expected_class)
             self.assertEqual(bx.shape, (2, 4))
             self.assertEqual(by.shape, (2, 2))
+
+    @pytest.mark.skipif(
+        not backend.SUPPORTS_RAGGED_TENSORS,
+        reason="Backend does not support ragged tensors",
+    )
+    def test_ragged_tensors(self):
+        x = tf.ragged.constant(
+            [[[0.0, 1.0]], [[2.0, 3.0], [4.0, 5.0]]], ragged_rank=1
+        )
+        y = tf.ragged.constant(
+            [[[0.0, 1.0]], [[0.0, 1.0], [0.0, 1.0]]], ragged_rank=1
+        )
+
+        def generate():
+            for _ in range(4):
+                yield x, y
+
+        adapter = generator_data_adapter.GeneratorDataAdapter(generate())
+
+        if backend.backend() == "tensorflow":
+            it = adapter.get_tf_dataset()
+            expected_class = tf.RaggedTensor
+
+        for batch in it:
+            self.assertEqual(len(batch), 2)
+            bx, by = batch
+            self.assertIsInstance(bx, expected_class)
+            self.assertIsInstance(by, expected_class)
+            self.assertEqual(bx.shape, (2, None, 2))
+            self.assertEqual(by.shape, (2, None, 2))