Add tf.RaggedTensor support to Embedding layer. (#21171)

Adds support for indices indices in the form of a `tf.RaggedTensor` to the `Embedding` layer by adding support to `ops.take`. The output is also ragged.

Also:
- adds support for negative indices in the sparse tensor use case.
- adds support for ragged tensors in `TestCase.run_layer_test`.

Author: hertschuh

keras/src/backend/tensorflow/numpy.py

@@ -2185,42 +2185,50 @@ def swapaxes(x, axis1, axis2):
 
 
 def take(x, indices, axis=None):
+    x = convert_to_tensor(x)
+    if axis is None:
+        x = tf.reshape(x, (-1,))
+        axis = 0
+
+    def fix_negative_indices(i):
+        # Correct the indices using "fill" mode which is the same as in jax
+        return tf.where(i < 0, i + tf.cast(tf.shape(x)[axis], i.dtype), i)
+
     if isinstance(indices, tf.SparseTensor):
         if x.dtype not in (tf.float16, tf.float32, tf.float64, tf.bfloat16):
             warnings.warn(
                 "`take` with the TensorFlow backend does not support "
                 f"`x.dtype={x.dtype}` when `indices` is a sparse tensor; "
                 "densifying `indices`."
             )
-            return take(x, convert_to_tensor(indices, sparse=False), axis=axis)
-        if axis is None:
-            x = tf.reshape(x, (-1,))
+            indices = convert_to_tensor(indices, sparse=False)
         elif axis != 0:
             warnings.warn(
                 "`take` with the TensorFlow backend does not support "
                 f"`axis={axis}` when `indices` is a sparse tensor; "
                 "densifying `indices`."
             )
-            return take(x, convert_to_tensor(indices, sparse=False), axis=axis)
-        output = tf.nn.safe_embedding_lookup_sparse(
-            embedding_weights=tf.convert_to_tensor(x),
-            sparse_ids=tf.sparse.expand_dims(indices, axis=-1),
-            default_id=0,
-        )
-        output.set_shape(indices.shape + output.shape[len(indices.shape) :])
-        return output
+            indices = convert_to_tensor(indices, sparse=False)
+        else:
+            indices = sparse.sparse_with_values(
+                indices, fix_negative_indices(indices.values)
+            )
+            # `expand_dims` on `indices` prevents combiner from being applied.
+            output = tf.nn.safe_embedding_lookup_sparse(
+                embedding_weights=tf.convert_to_tensor(x),
+                sparse_ids=tf.sparse.expand_dims(indices, axis=-1),
+                default_id=0,
+            )
+            output.set_shape(indices.shape + output.shape[len(indices.shape) :])
+            return output
+    elif isinstance(indices, tf.RaggedTensor):
+        indices = indices.with_values(fix_negative_indices(indices.values))
+        if axis == 0:
+            return tf.nn.embedding_lookup(x, indices)
+        else:
+            return tf.gather(x, indices, axis=axis)
 
-    x = convert_to_tensor(x)
-    indices = convert_to_tensor(indices)
-    if axis is None:
-        x = tf.reshape(x, [-1])
-        axis = 0
-    # Correct the indices using "fill" mode which is the same as in jax
-    indices = tf.where(
-        indices < 0,
-        indices + tf.cast(tf.shape(x)[axis], indices.dtype),
-        indices,
-    )
+    indices = fix_negative_indices(convert_to_tensor(indices))
     return tf.gather(x, indices, axis=axis)
 
 

keras/src/layers/core/embedding.py

@@ -8,6 +8,7 @@
 from keras.src import quantizers
 from keras.src import regularizers
 from keras.src.api_export import keras_export
+from keras.src.backend import KerasTensor
 from keras.src.layers.layer import Layer
 
 
@@ -156,6 +157,12 @@ def compute_mask(self, inputs, mask=None):
     def compute_output_shape(self, input_shape):
         return (*input_shape, self.output_dim)
 
+    def compute_output_spec(self, inputs):
+        output_shape = (*inputs.shape, self.output_dim)
+        return KerasTensor(
+            output_shape, dtype=self.compute_dtype, ragged=inputs.ragged
+        )
+
     def enable_lora(
         self,
         rank,

keras/src/layers/core/embedding_test.py

@@ -61,6 +61,27 @@ def test_sparse(self):
             supports_masking=False,
         )
 
+    @pytest.mark.skipif(
+        not backend.SUPPORTS_RAGGED_TENSORS,
+        reason="Backend does not support ragged tensors.",
+    )
+    def test_ragged(self):
+        self.run_layer_test(
+            layers.Embedding,
+            {"input_dim": 5, "output_dim": 4},
+            input_shape=(2, 3),
+            input_dtype="int32",
+            input_ragged=True,
+            expected_output_shape=(2, None, 4),
+            expected_output_ragged=True,
+            expected_num_trainable_weights=1,
+            expected_num_non_trainable_weights=0,
+            expected_num_seed_generators=0,
+            expected_num_losses=0,
+            supports_masking=False,
+            # run_training_check=False,
+        )
+
     def test_correctness(self):
         layer = layers.Embedding(input_dim=3, output_dim=2)
         layer.build()

keras/src/ops/numpy.py

@@ -5425,15 +5425,17 @@ def compute_output_spec(self, x, indices):
         x_shape = list(x.shape)
         if isinstance(indices, KerasTensor):
             indices_shape = list(indices.shape)
+            ragged = indices.ragged
         else:
             indices_shape = list(getattr(np.array(indices), "shape", []))
+            ragged = False
         if self.axis is None:
             return KerasTensor(indices_shape, dtype=x.dtype)
 
         # make sure axis is non-negative
         axis = len(x_shape) + self.axis if self.axis < 0 else self.axis
         output_shape = x_shape[:axis] + indices_shape + x_shape[axis + 1 :]
-        return KerasTensor(output_shape, dtype=x.dtype)
+        return KerasTensor(output_shape, dtype=x.dtype, ragged=ragged)
 
 
 @keras_export(["keras.ops.take", "keras.ops.numpy.take"])

keras/src/ops/numpy_test.py

@@ -3078,17 +3078,61 @@ def test_take_sparse(self, dtype, axis):
         if backend.backend() == "tensorflow":
             import tensorflow as tf
 
-            indices = tf.SparseTensor([[0, 0], [1, 2]], [1, 2], (2, 3))
+            indices = tf.SparseTensor([[0, 0], [1, 2]], [-1, 2], (2, 3))
         elif backend.backend() == "jax":
             import jax.experimental.sparse as jax_sparse
 
-            indices = jax_sparse.BCOO(([1, 2], [[0, 0], [1, 2]]), shape=(2, 3))
+            indices = jax_sparse.BCOO(([-1, 2], [[0, 0], [1, 2]]), shape=(2, 3))
 
         self.assertAllClose(
             knp.take(x, indices, axis=axis),
             np.take(x, backend.convert_to_numpy(indices), axis=axis),
         )
 
+    @parameterized.named_parameters(
+        named_product(
+            [
+                {"testcase_name": "axis_none", "axis": None},
+                {"testcase_name": "axis_0", "axis": 0},
+                {"testcase_name": "axis_1", "axis": 1},
+                {"testcase_name": "axis_minus1", "axis": -1},
+            ],
+            dtype=[
+                "float16",
+                "float32",
+                "float64",
+                "uint8",
+                "int8",
+                "int16",
+                "int32",
+            ],
+        )
+    )
+    @pytest.mark.skipif(
+        not backend.SUPPORTS_RAGGED_TENSORS,
+        reason="Backend does not support ragged tensors.",
+    )
+    def test_take_ragged(self, dtype, axis):
+        rng = np.random.default_rng(0)
+        x = (4 * rng.standard_normal((3, 4, 5))).astype(dtype)
+
+        if backend.backend() == "tensorflow":
+            import tensorflow as tf
+
+            indices = tf.ragged.constant([[2], [0, -1, 1]])
+            mask = backend.convert_to_numpy(tf.ones_like(indices))
+
+        if axis == 0:
+            mask = np.expand_dims(mask, (2, 3))
+        elif axis == 1:
+            mask = np.expand_dims(mask, (2,))
+
+        self.assertAllClose(
+            knp.take(x, indices, axis=axis),
+            np.take(x, backend.convert_to_numpy(indices), axis=axis)
+            * mask.astype(dtype),
+        )
+
     def test_take_along_axis(self):
         x = np.arange(24).reshape([1, 2, 3, 4])
         indices = np.ones([1, 4, 1, 1], dtype=np.int32)

keras/src/testing/test_case.py

@@ -16,6 +16,7 @@
 from keras.src.backend.common import standardize_dtype
 from keras.src.backend.common.global_state import clear_session
 from keras.src.backend.common.keras_tensor import KerasTensor
+from keras.src.losses.losses import MeanSquaredError
 from keras.src.models import Model
 from keras.src.utils import traceback_utils
 
@@ -100,6 +101,22 @@ def assertSparse(self, x, sparse=True):
                 f"Backend {backend.backend()} does not support sparse tensors",
             )
 
+    def assertRagged(self, x, ragged=True):
+        if isinstance(x, KerasTensor):
+            self.assertEqual(x.ragged, ragged)
+        elif backend.backend() == "tensorflow":
+            import tensorflow as tf
+
+            if ragged:
+                self.assertIsInstance(x, tf.RaggedTensor)
+            else:
+                self.assertNotIsInstance(x, tf.RaggedTensor)
+        else:
+            self.assertFalse(
+                ragged,
+                f"Backend {backend.backend()} does not support ragged tensors",
+            )
+
     def assertDType(self, x, dtype, msg=None):
         if hasattr(x, "dtype"):
             x_dtype = backend.standardize_dtype(x.dtype)
@@ -159,11 +176,13 @@ def run_layer_test(
         input_shape=None,
         input_dtype=None,
         input_sparse=False,
+        input_ragged=False,
         input_data=None,
         call_kwargs=None,
         expected_output_shape=None,
         expected_output_dtype=None,
         expected_output_sparse=False,
+        expected_output_ragged=False,
         expected_output=None,
         expected_num_trainable_weights=None,
         expected_num_non_trainable_weights=None,
@@ -188,6 +207,8 @@ def run_layer_test(
             input_dtype: Corresponding input dtype.
             input_sparse: Whether the input is a sparse tensor (this requires
                 the backend to support sparse tensors).
+            input_ragged: Whether the input is a ragged tensor (this requires
+                the backend to support ragged tensors).
             input_data: Tensor (or list/dict of tensors)
                 to call the layer on.
             call_kwargs: Dict of arguments to use when calling the
@@ -198,6 +219,8 @@ def run_layer_test(
             expected_output_dtype: dtype expected as output.
             expected_output_sparse: Whether the output is expected to be sparse
                 (this requires the backend to support sparse tensors).
+            expected_output_ragged: Whether the output is expected to be ragged
+                (this requires the backend to support ragged tensors).
             expected_output: Expected output tensor -- only
                 to be specified if input_data is provided.
             expected_num_trainable_weights: Expected number
@@ -280,7 +303,7 @@ def run_layer_test(
         if input_data is not None or input_shape is not None:
             if input_data is None:
                 input_data = create_eager_tensors(
-                    input_shape, input_dtype, input_sparse
+                    input_shape, input_dtype, input_sparse, input_ragged
                 )
             layer = layer_cls(**init_kwargs)
             if isinstance(input_data, dict):
@@ -357,7 +380,13 @@ def run_output_asserts(layer, output, eager=False):
             if expected_output_shape is not None:
 
                 def verify_shape(expected_shape, x):
-                    return expected_shape == x.shape
+                    shape = x.shape
+                    if len(shape) != len(expected_shape):
+                        return False
+                    for expected_dim, dim in zip(expected_shape, shape):
+                        if expected_dim is not None and expected_dim != dim:
+                            return False
+                    return True
 
                 shapes_match = tree.map_structure_up_to(
                     output, verify_shape, expected_output_shape, output
@@ -383,6 +412,9 @@ def verify_dtype(expected_dtype, x):
             if expected_output_sparse:
                 for x in tree.flatten(output):
                     self.assertSparse(x)
+            if expected_output_ragged:
+                for x in tree.flatten(output):
+                    self.assertRagged(x)
             if eager:
                 if expected_output is not None:
                     self.assertEqual(type(expected_output), type(output))
@@ -430,7 +462,11 @@ def data_generator():
             jit_compile = "auto"
             if backend.backend() == "tensorflow" and input_sparse:
                 jit_compile = False
-            model.compile(optimizer="sgd", loss="mse", jit_compile=jit_compile)
+            model.compile(
+                optimizer="sgd",
+                loss=MeanSquaredError(reduction="sum"),
+                jit_compile=jit_compile,
+            )
             model.fit(data_generator(), steps_per_epoch=1, verbose=0)
 
         # Build test.
@@ -452,13 +488,13 @@ def data_generator():
             if input_shape is None:
                 keras_tensor_inputs = tree.map_structure(
                     lambda x: create_keras_tensors(
-                        ops.shape(x), x.dtype, input_sparse
+                        ops.shape(x), x.dtype, input_sparse, input_ragged
                     ),
                     input_data,
                 )
             else:
                 keras_tensor_inputs = create_keras_tensors(
-                    input_shape, input_dtype, input_sparse
+                    input_shape, input_dtype, input_sparse, input_ragged
                 )
             layer = layer_cls(**init_kwargs)
             if isinstance(keras_tensor_inputs, dict):
@@ -589,22 +625,24 @@ def uses_cpu():
     return False
 
 
-def create_keras_tensors(input_shape, dtype, sparse):
+def create_keras_tensors(input_shape, dtype, sparse, ragged):
     if isinstance(input_shape, dict):
         return {
             utils.removesuffix(k, "_shape"): KerasTensor(
-                v, dtype=dtype[k], sparse=sparse
+                v, dtype=dtype[k], sparse=sparse, ragged=ragged
             )
             for k, v in input_shape.items()
         }
     return map_shape_dtype_structure(
-        lambda shape, dt: KerasTensor(shape, dtype=dt, sparse=sparse),
+        lambda shape, dt: KerasTensor(
+            shape, dtype=dt, sparse=sparse, ragged=ragged
+        ),
         input_shape,
         dtype,
     )
 
 
-def create_eager_tensors(input_shape, dtype, sparse):
+def create_eager_tensors(input_shape, dtype, sparse, ragged):
     from keras.src.backend import random
 
     if set(tree.flatten(dtype)).difference(
@@ -651,6 +689,21 @@ def create_fn(shape, dt):
                 f"Sparse is unsupported with backend {backend.backend()}"
             )
 
+    elif ragged:
+        if backend.backend() == "tensorflow":
+            import tensorflow as tf
+
+            def create_fn(shape, dt):
+                rng = np.random.default_rng(0)
+                x = (4 * rng.standard_normal(shape)).astype(dt)
+                x = np.multiply(x, rng.random(shape) < 0.7)
+                return tf.RaggedTensor.from_tensor(x, padding=0)
+
+        else:
+            raise ValueError(
+                f"Ragged is unsupported with backend {backend.backend()}"
+            )
+
     else:
 
         def create_fn(shape, dt):