[OpenVINO Backend] support ops.slice_update

keras/src/backend/openvino/core.py

@@ -810,10 +810,117 @@ def prepare_slice_index(val):
 
 
 def slice_update(inputs, start_indices, updates):
-    raise NotImplementedError(
-        "`slice_update` is not supported with openvino backend"
+    inputs = get_ov_output(inputs)
+    updates_tensor = get_ov_output(updates)
+
+    if isinstance(start_indices, (list, np.ndarray)):
+        start_indices = tuple(start_indices)
+    assert isinstance(start_indices, tuple), (
+        "`slice_update` is not supported by openvino backend"
+        " for `start_indices` of type {}".format(type(start_indices))
     )
 
+    zero_scalar = ov_opset.constant(0, Type.i32)
+    one_scalar = ov_opset.constant(1, Type.i32)
+    zero_tensor = ov_opset.constant([0], Type.i32)
+    one_tensor = ov_opset.constant([1], Type.i32)
+
+    processed_start_indices = []
+    for idx in start_indices:
+        val = get_ov_output(idx)
+        if not val.get_element_type().is_integral():
+            raise ValueError("`slice_update` requires integral start_indices")
+        if val.get_element_type() != Type.i32:
+            val = ov_opset.convert(val, Type.i32).output(0)
+        if val.get_partial_shape().rank.get_length() == 0:
+            val = ov_opset.unsqueeze(val, zero_scalar).output(0)
+        processed_start_indices.append(val)
+
+    updates_shape = ov_opset.shape_of(updates_tensor, Type.i32).output(0)
+    rank = updates_tensor.get_partial_shape().rank.get_length()
+
+    total_elements = ov_opset.reduce_prod(
+        updates_shape, zero_tensor, keep_dims=False
+    ).output(0)
+
+    # Create a single range for all indices
+    flat_indices = ov_opset.range(
+        zero_scalar, total_elements, one_scalar, output_type=Type.i32
+    ).output(0)
+
+    dim_sizes = []
+    strides = []
+
+    for dim in range(rank):
+        dim_size = ov_opset.gather(
+            updates_shape, ov_opset.constant([dim], Type.i32), zero_scalar
+        ).output(0)
+        dim_size_scalar = ov_opset.squeeze(dim_size, zero_tensor).output(0)
+        dim_sizes.append(dim_size_scalar)
+
+        # Compute stride (product of dimensions after current)
+        if dim < rank - 1:
+            remaining_dims = ov_opset.slice(
+                updates_shape,
+                ov_opset.constant([dim + 1], Type.i32),
+                ov_opset.constant([rank], Type.i32),
+                one_tensor,
+                zero_tensor,
+            ).output(0)
+            stride = ov_opset.reduce_prod(
+                remaining_dims, zero_tensor, keep_dims=False
+            ).output(0)
+        else:
+            stride = one_scalar
+        strides.append(stride)
+
+    coord_tensors = []
+    for dim in range(rank):
+        # Calculate coordinates for this dimension
+        coords = ov_opset.mod(
+            ov_opset.divide(flat_indices, strides[dim]).output(0),
+            dim_sizes[dim],
+        ).output(0)
+        coord_tensors.append(coords)
+
+    coord_tensors_unsqueezed = []
+    for coord in coord_tensors:
+        coord_unsqueezed = ov_opset.unsqueeze(coord, one_tensor).output(0)
+        coord_tensors_unsqueezed.append(coord_unsqueezed)
+
+    # Create index matrix [total_elements, rank] in one operation
+    indices_matrix = ov_opset.concat(coord_tensors_unsqueezed, axis=1).output(0)
+
+    # Broadcast start indices
+    start_tensor = ov_opset.concat(processed_start_indices, axis=0).output(0)
+    start_reshaped = ov_opset.reshape(
+        start_tensor, ov_opset.constant([1, rank], Type.i32), special_zero=False
+    ).output(0)
+
+    # Broadcast to match indices matrix shape
+    broadcast_shape = ov_opset.concat(
+        [
+            ov_opset.unsqueeze(total_elements, zero_tensor).output(0),
+            one_tensor,
+        ],
+        axis=0,
+    ).output(0)
+
+    start_broadcast = ov_opset.tile(start_reshaped, broadcast_shape).output(0)
+
+    # Add offset to get absolute indices
+    absolute_indices = ov_opset.add(indices_matrix, start_broadcast).output(0)
+
+    updates_flat = ov_opset.reshape(
+        updates_tensor,
+        ov_opset.unsqueeze(total_elements, zero_tensor).output(0),
+        special_zero=False,
+    ).output(0)
+    result = ov_opset.scatter_nd_update(
+        inputs, absolute_indices, updates_flat
+    ).output(0)
+    return OpenVINOKerasTensor(result)
+
 
 def while_loop(
     cond,

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -164,17 +164,14 @@ CoreOpsCallsTests::test_map_basic_call
 CoreOpsCallsTests::test_scan_basic_call
 CoreOpsCallsTests::test_scatter_basic_call
 CoreOpsCallsTests::test_scatter_update_basic_call
-CoreOpsCallsTests::test_slice_update_basic_call
 CoreOpsCallsTests::test_switch_basic_call
 CoreOpsCallsTests::test_unstack_basic_functionality
 CoreOpsCorrectnessTest::test_associative_scan
 CoreOpsCorrectnessTest::test_cond
-CoreOpsCorrectnessTest::test_dynamic_slice
 CoreOpsCorrectnessTest::test_fori_loop
 CoreOpsCorrectnessTest::test_map
 CoreOpsCorrectnessTest::test_scan
 CoreOpsCorrectnessTest::test_scatter
-CoreOpsCorrectnessTest::test_slice_update
 CoreOpsCorrectnessTest::test_switch 
 CoreOpsCorrectnessTest::test_unstack
 CoreOpsCorrectnessTest::test_vectorized_map