Fix Embedding test with ragged tensors on GPU. (#21177)

The loss needs to not have any non-compilable op.

Author: hertschuh

keras/src/testing/test_case.py

@@ -16,7 +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.losses.loss import Loss
 from keras.src.models import Model
 from keras.src.utils import traceback_utils
 
@@ -446,6 +446,11 @@ def data_generator():
                 while True:
                     yield data
 
+            # Single op loss to avoid compilation issues with ragged / sparse.
+            class TestLoss(Loss):
+                def __call__(self, y_true, y_pred, sample_weight=None):
+                    return ops.sum(y_pred)
+
             # test the "default" path for each backend by setting
             # jit_compile="auto".
             # for tensorflow and jax backends auto is jitted
@@ -463,9 +468,7 @@ def data_generator():
             if backend.backend() == "tensorflow" and input_sparse:
                 jit_compile = False
             model.compile(
-                optimizer="sgd",
-                loss=MeanSquaredError(reduction="sum"),
-                jit_compile=jit_compile,
+                optimizer="sgd", loss=TestLoss(), jit_compile=jit_compile
             )
             model.fit(data_generator(), steps_per_epoch=1, verbose=0)