[OpenVINO Backend] update __getitem__

Author: Mohamed-Ashraf273

keras/src/backend/openvino/core.py

@@ -1,3 +1,4 @@
+import builtins
 import contextlib
 import warnings
 
@@ -308,30 +309,94 @@ def __ne__(self, other):
         return OpenVINOKerasTensor(ov_opset.not_equal(first, other).output(0))
 
     def __getitem__(self, indices):
-        # now it has limited functionaly
-        # and supports only a case with one integer index in indices
-        # other indices must be None
         data = self.output
-        axis = []
-        gather_index = None
-        if isinstance(indices, int):
+        axes, gather_indices_nodes = [], []
+        slice_axes, slice_starts, slice_ends, slice_steps = [], [], [], []
+        unsqueeze_axes, ellipsis_index = [], None
+
+        if not isinstance(indices, tuple):
             indices = (indices,)
-        assert isinstance(indices, tuple), "only tuple is supported"
+
         for dim, index in enumerate(indices):
             if isinstance(index, int):
-                axis.append(dim)
-                gather_index = ov_opset.constant(index, Type.i32)
-            else:
-                assert (
+                axes.append(dim)
+                index_node = ov_opset.constant([index], Type.i32).output(0)
+                gather_indices_nodes.append(index_node)
+            elif isinstance(index, builtins.slice):
+                if (
                     index.start is None
                     and index.stop is None
                     and index.step is None
+                ):
+                    continue
+                slice_axes.append(dim)
+                slice_starts.append(0 if index.start is None else index.start)
+                slice_ends.append(
+                    2**31 - 1 if index.stop is None else index.stop
                 )
-        assert len(axis) == 1, "axis must contain one element"
-        axis = ov_opset.constant(axis, Type.i32)
-        return OpenVINOKerasTensor(
-            ov_opset.gather(data, gather_index, axis).output(0)
-        )
+                slice_steps.append(1 if index.step is None else index.step)
+            elif index is Ellipsis:
+                ellipsis_index = dim
+            elif index is None:
+                unsqueeze_axes.append(dim)
+            elif isinstance(index, OpenVINOKerasTensor):
+                axes.append(dim)
+                gather_indices_nodes.append(get_ov_output(index))
+            else:
+                raise ValueError(
+                    f"""Unsupported index type {type(index)}
+                    in OpenVINOKerasTensor.__getitem__"""
+                )
+
+        if slice_axes:
+            step = ov_opset.constant(slice_steps, Type.i32).output(0)
+            start = ov_opset.constant(slice_starts, Type.i32).output(0)
+            stop = ov_opset.constant(slice_ends, Type.i32).output(0)
+            axes_const = ov_opset.constant(slice_axes, Type.i32).output(0)
+            data = ov_opset.slice(data, start, stop, step, axes_const).output(0)
+
+        if axes:
+            gather_indices_const = (
+                gather_indices_nodes[0]
+                if len(gather_indices_nodes) == 1
+                else ov_opset.concat(gather_indices_nodes, axis=0).output(0)
+            )
+            if len(axes) == 1:
+                data = ov_opset.gather(
+                    data, gather_indices_const, axes[0]
+                ).output(0)
+                data = ov_opset.squeeze(data, axes[0]).output(0)
+            else:
+                adjusted_axes = [
+                    ax - sum(1 for unsq in unsqueeze_axes if unsq <= ax)
+                    for ax in axes
+                ]
+                rank = len(data.get_partial_shape())
+                remaining_axes = [
+                    i for i in range(rank) if i not in adjusted_axes
+                ]
+                perm = ov_opset.constant(
+                    adjusted_axes + remaining_axes, Type.i32
+                )
+                data = ov_opset.transpose(data, perm).output(0)
+                data = ov_opset.gather_nd(data, gather_indices_const).output(0)
+
+        if unsqueeze_axes:
+            expanded_rank = data.get_partial_shape().rank.get_length() + len(
+                unsqueeze_axes
+            )
+            adjusted_unsqueeze = []
+            for ax in unsqueeze_axes:
+                ax -= sum(1 for s in axes if s < ax)
+                if ellipsis_index is not None and ax > ellipsis_index:
+                    ax += expanded_rank - len(indices)
+                adjusted_unsqueeze.append(ax)
+            unsqueeze_const = ov_opset.constant(
+                adjusted_unsqueeze, Type.i32
+            ).output(0)
+            data = ov_opset.unsqueeze(data, unsqueeze_const).output(0)
+
+        return OpenVINOKerasTensor(data)
 
     def __len__(self):
         ov_output = self.output

keras/src/backend/openvino/excluded_tests.txt

@@ -1,6 +1,5 @@
 keras/src/activations
 keras/src/backend/common/dtypes_test.py
-keras/src/backend/common/variables_test.py
 keras/src/callbacks/early_stopping_test.py
 keras/src/dtype_policies/dtype_policy_map_test.py
 keras/src/layers/attention