[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,119 @@ 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):
+        rank = len(data.get_partial_shape())
+        axes, gather_indices_nodes = [], []
+        slice_axes, slice_starts, slice_ends, slice_steps = [], [], [], []
+        unsqueeze_axes = []
+
+        if not isinstance(indices, tuple):
             indices = (indices,)
-        assert isinstance(indices, tuple), "only tuple is supported"
+
+        if any(i is Ellipsis for i in indices):
+            ellipsis_pos = indices.index(Ellipsis)
+            num_specified = sum(
+                i is not Ellipsis and i is not None for i in indices
+            )
+            num_missing = rank - num_specified
+            indices = (
+                indices[:ellipsis_pos]
+                + (builtins.slice(None),) * num_missing
+                + indices[ellipsis_pos + 1 :]
+            )
+
+        def count_unsqueeze_before(dim):
+            return sum(1 for i in range(dim) if indices[i] is None)
+
+        partial_shape = data.get_partial_shape()
+
         for dim, index in enumerate(indices):
             if isinstance(index, int):
-                axis.append(dim)
-                gather_index = ov_opset.constant(index, Type.i32)
+                actual_dim = dim - count_unsqueeze_before(dim)
+                if not (0 <= actual_dim < rank):
+                    raise IndexError(
+                        f"""Index {index} is out of bounds for
+                        axis {dim} with rank {rank}"""
+                    )
+                length = partial_shape[actual_dim].get_length()
+                idx_value = index if index >= 0 else length + index
+                axes.append(dim)
+                gather_indices_nodes.append(
+                    ov_opset.constant([idx_value], Type.i32).output(0)
+                )
+            elif isinstance(index, builtins.slice):
+                if index == builtins.slice(None):
+                    continue
+                if index.step is not None and index.step < 0:
+                    raise ValueError("OpenVINO doesn't support negative steps")
+                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
+                )
+                slice_steps.append(1 if index.step is None else index.step)
+            elif index is None:
+                unsqueeze_axes.append(dim)
+            elif isinstance(index, OpenVINOKerasTensor):
+                axes.append(dim)
+                gather_indices_nodes.append(get_ov_output(index))
             else:
-                assert (
-                    index.start is None
-                    and index.stop is None
-                    and index.step is None
+                raise ValueError(
+                    f"""Unsupported index type {type(index)}
+                    in OpenVINOKerasTensor.__getitem__"""
                 )
-        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)
-        )
+
+        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)
+            adjusted_slice_axes = [
+                ax - sum(1 for unsq in unsqueeze_axes if unsq <= ax)
+                for ax in slice_axes
+            ]
+            axes_const = ov_opset.constant(
+                adjusted_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)
+            )
+            adjusted_axes = [
+                ax - sum(1 for unsq in unsqueeze_axes if unsq <= ax)
+                for ax in axes
+            ]
+            if len(axes) == 1:
+                data = ov_opset.gather(
+                    data, gather_indices_const, adjusted_axes[0]
+                ).output(0)
+                data = ov_opset.squeeze(data, adjusted_axes[0]).output(0)
+            else:
+                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:
+            adjusted_unsqueeze = []
+            for ax in unsqueeze_axes:
+                ax -= sum(1 for s in axes if s < ax)
+                ax -= sum(1 for s in slice_axes if s < ax)
+                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