[OpenVINO Backend] support ops.split (#21296)

Author: Mohamed-Ashraf273

keras/src/backend/openvino/core.py

@@ -450,8 +450,13 @@ def convert_to_tensor(x, dtype=None, sparse=None, ragged=None):
 def convert_to_numpy(x):
     if isinstance(x, np.ndarray):
         return x
-    elif isinstance(x, (int, float, list, tuple)):
+    elif isinstance(x, (int, float)):
         return np.array(x)
+    elif isinstance(x, (list, tuple)):
+        x_new = []
+        for elem in x:
+            x_new.append(convert_to_numpy(elem))
+        return np.array(x_new)
     elif np.isscalar(x):
         return x
     elif isinstance(x, ov.Tensor):

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -48,7 +48,6 @@ NumpyDtypeTest::test_round
 NumpyDtypeTest::test_searchsorted
 NumpyDtypeTest::test_signbit
 NumpyDtypeTest::test_sort
-NumpyDtypeTest::test_split
 NumpyDtypeTest::test_sqrt
 NumpyDtypeTest::test_std
 NumpyDtypeTest::test_subtract
@@ -112,7 +111,6 @@ NumpyOneInputOpsCorrectnessTest::test_signbit
 NumpyOneInputOpsCorrectnessTest::test_size
 NumpyOneInputOpsCorrectnessTest::test_slogdet
 NumpyOneInputOpsCorrectnessTest::test_sort
-NumpyOneInputOpsCorrectnessTest::test_split
 NumpyOneInputOpsCorrectnessTest::test_sqrt_int32
 NumpyOneInputOpsCorrectnessTest::test_squeeze
 NumpyOneInputOpsCorrectnessTest::test_std

keras/src/backend/openvino/numpy.py

@@ -1340,7 +1340,49 @@ def sort(x, axis=-1):
 
 
 def split(x, indices_or_sections, axis=0):
-    raise NotImplementedError("`split` is not supported with openvino backend")
+    x = get_ov_output(x)
+    axis_tensor = ov_opset.constant(axis, dtype=Type.i32).output(0)
+
+    shape_tensor = ov_opset.shape_of(x)
+    axis_i32 = ov_opset.constant([axis], dtype=Type.i32)
+    dim_at_axis_tensor = ov_opset.gather(
+        shape_tensor, axis_i32, ov_opset.constant(0, dtype=Type.i32)
+    )
+
+    if isinstance(indices_or_sections, int):
+        num_splits = indices_or_sections
+        splits = ov_opset.split(x, axis_tensor, num_splits=num_splits)
+        result = []
+        for i in range(num_splits):
+            result.append(OpenVINOKerasTensor(splits.output(i)))
+        return result
+
+    if isinstance(indices_or_sections, (list, tuple, np.ndarray)):
+        indices = list(indices_or_sections)
+        split_lengths = []
+        split_lengths.append(indices[0])
+        for i in range(1, len(indices)):
+            split_lengths.append(indices[i] - indices[i - 1])
+
+        last_index_tensor = ov_opset.constant(indices[-1], dtype=Type.i64)
+        remaining_length_tensor = ov_opset.subtract(
+            dim_at_axis_tensor, last_index_tensor
+        )
+
+        length_parts = []
+        length_parts.append(ov_opset.constant(split_lengths, dtype=Type.i64))
+        length_parts.append(remaining_length_tensor)
+        length_tensor = ov_opset.concat(length_parts, axis=0)
+
+        splits = ov_opset.variadic_split(x, axis_tensor, length_tensor)
+        result = []
+        for i in range(len(split_lengths) + 1):
+            result.append(OpenVINOKerasTensor(splits.output(i)))
+        return result
+
+    raise TypeError(
+        f"unsupported type of indices_or_sections: {type(indices_or_sections)}"
+    )
 
 
 def stack(x, axis=0):