Support shape -1 for slice op in the jax backend (#21501)

* Support shape -1 for slice op in the jax backend.

* Fix comments.

* Update names

* Address Gemini reviews.

Author: SamanehSaadat

keras/src/backend/jax/core.py

@@ -300,7 +300,13 @@ def scatter_update(inputs, indices, updates):
 
 
 def slice(inputs, start_indices, shape):
-    return jax.lax.dynamic_slice(inputs, start_indices, shape)
+    # If shape[i] is -1, all remaining elements in dimension i are included in
+    # the slice.
+    final_shape = tuple(
+        inputs.shape[i] - start_indices[i] if s == -1 else s
+        for i, s in enumerate(shape)
+    )
+    return jax.lax.dynamic_slice(inputs, start_indices, final_shape)
 
 
 def slice_update(inputs, start_indices, updates):

keras/src/ops/core.py

@@ -397,7 +397,20 @@ def call(self, inputs, start_indices):
         return backend.core.slice(inputs, start_indices, self.shape)
 
     def compute_output_spec(self, inputs, start_indices):
-        return KerasTensor(self.shape, dtype=inputs.dtype)
+        if any(s == -1 for s in self.shape) and isinstance(
+            start_indices, KerasTensor
+        ):
+            raise ValueError(
+                "When using -1 in `shape`, `start_indices` should not be a "
+                "KerasTensor. "
+            )
+        # If self.shape[i] is -1, all remaining elements in dimension i are
+        # included in the slice.
+        final_shape = tuple(
+            inputs.shape[i] - start_indices[i] if s == -1 else s
+            for i, s in enumerate(self.shape)
+        )
+        return KerasTensor(final_shape, dtype=inputs.dtype)
 
 
 @keras_export("keras.ops.slice")

keras/src/ops/core_test.py

@@ -265,6 +265,19 @@ def test_slice(self):
         shape = (2, 2)
         self.assertEqual(core.slice(inputs, start_indices, shape).shape, (2, 2))
 
+    def test_slice_negative_one_shape(self):
+        inputs = KerasTensor(shape=(3, 3), dtype="float32")
+        start_indices = (1, 1)
+        shape = (-1, -1)
+        self.assertEqual(core.slice(inputs, start_indices, shape).shape, (2, 2))
+
+    def test_slice_negative_one_shape_raises(self):
+        inputs = KerasTensor(shape=(3, 3), dtype="float32")
+        start_indices = KerasTensor(shape=(2,), dtype="int32")
+        shape = (-1, -1)
+        with self.assertRaises(ValueError):
+            core.slice(inputs, start_indices, shape)
+
     def test_slice_update(self):
         inputs = KerasTensor((4, 4))
         start_indices = KerasTensor((2,))