[OpenVNIO Backend] update __getitem__

Author: Mohamed-Ashraf273

keras/src/backend/openvino/core.py

@@ -333,7 +333,8 @@ def __getitem__(self, indices):
         def count_unsqueeze_before(dim):
             return sum(1 for i in range(dim) if indices[i] is None)
 
-        partial_shape = data.get_partial_shape()
+        partial_shape = ov_opset.shape_of(data, Type.i32)
+        zero_const = ov_opset.constant(0, Type.i32)
 
         for dim, index in enumerate(indices):
             if isinstance(index, bool):
@@ -347,12 +348,19 @@ def count_unsqueeze_before(dim):
                         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)
+                length = ov_opset.gather(
+                    partial_shape,
+                    ov_opset.constant([actual_dim], Type.i32),
+                    zero_const,
                 )
+                if index >= 0:
+                    idx_value = ov_opset.constant([index], Type.i32)
+                else:
+                    idx_value = ov_opset.add(
+                        ov_opset.constant([index], Type.i32), length
+                    )
+                axes.append(dim)
+                gather_indices_nodes.append(idx_value.output(0))
             elif isinstance(index, builtins.slice):
                 if index == builtins.slice(None):
                     continue
@@ -381,25 +389,18 @@ def count_unsqueeze_before(dim):
                         "OpenVINO backend does not "
                         "support multi-dimensional indexing"
                     )
+                if len(index_shape) == 0:
+                    index = ov_opset.unsqueeze(index, zero_const).output(0)
                 if index_type != Type.i32:
                     index = ov_opset.convert(index, Type.i32).output(0)
-                shape_tensor = ov_opset.shape_of(data)
+                shape_tensor = ov_opset.shape_of(data, Type.i32)
                 axis_i32 = ov_opset.constant([dim], dtype=Type.i32)
-                dim_size = ov_opset.gather(
-                    shape_tensor, axis_i32, ov_opset.constant(0, Type.i32)
-                )
-                dim_size = ov_opset.convert(dim_size, Type.i32)
-                zero = ov_opset.constant(0, Type.i32)
-                is_negative = ov_opset.less(index, zero)
+                dim_size = ov_opset.gather(shape_tensor, axis_i32, zero_const)
+                is_negative = ov_opset.less(index, zero_const)
                 adjusted_index = ov_opset.add(index, dim_size)
                 index = ov_opset.select(
                     is_negative, adjusted_index, index
                 ).output(0)
-                index_shape = index.get_partial_shape()
-                if len(index_shape) == 0:
-                    index = ov_opset.unsqueeze(
-                        index, ov_opset.constant(0, Type.i32)
-                    ).output(0)
                 gather_indices_nodes.append(index)
             else:
                 raise ValueError(

keras/src/ops/core_test.py

@@ -169,6 +169,172 @@ def test_convert_to_tensor(self):
 
 
 class CoreOpsCorrectnessTest(testing.TestCase):
+    def test_getitem(self):
+        self.np_tensor = np.arange(24).reshape(2, 3, 4)
+        self.tensor = ops.convert_to_tensor(self.np_tensor)
+
+        t = self.tensor[1]
+        n = self.np_tensor[1]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[1, 2, 3]
+        n = self.np_tensor[1, 2, 3]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[1:2]
+        n = self.np_tensor[1:2]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[1:2, 2:3, 3:4]
+        n = self.np_tensor[1:2, 2:3, 3:4]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[1:2, None]
+        n = self.np_tensor[1:2, None]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[1:2, 2:3, ...]
+        n = self.np_tensor[1:2, 2:3, ...]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[1:2, ..., 3:4]
+        n = self.np_tensor[1:2, ..., 3:4]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[None, ..., 3:4, None]
+        n = self.np_tensor[None, ..., 3:4, None]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[1:2:None]
+        n = self.np_tensor[1:2:None]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[:, 2]
+        n = self.np_tensor[:, 2]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[None]
+        n = self.np_tensor[None]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[None, None]
+        n = self.np_tensor[None, None]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[...]
+        n = self.np_tensor[...]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[..., 1]
+        n = self.np_tensor[..., 1]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[..., 1, 2]
+        n = self.np_tensor[..., 1, 2]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[..., -1, 2]
+        n = self.np_tensor[..., -1, 2]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[..., -1:-2, 2]
+        n = self.np_tensor[..., -1:-2, 2]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[..., None, None]
+        n = self.np_tensor[..., None, None]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[None, ..., None]
+        n = self.np_tensor[None, ..., None]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[1, 2, None, ..., None]
+        n = self.np_tensor[1, 2, None, ..., None]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[None, ..., 1, 2]
+        n = self.np_tensor[None, ..., 1, 2]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[1, None, 2]
+        n = self.np_tensor[1, None, 2]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        index_tensor = ops.convert_to_tensor(np.array(1, dtype=np.int32))
+        t = self.tensor[index_tensor]
+        n = self.np_tensor[ops.convert_to_numpy(index_tensor)]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        index_tensor = ops.convert_to_tensor(np.array(1, dtype=np.int32))
+        t = self.tensor[index_tensor, 2, None]
+        n = self.np_tensor[ops.convert_to_numpy(index_tensor), 2, None]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        index_tensor = ops.convert_to_tensor(np.array(-2, dtype=np.int32))
+        t = self.tensor[index_tensor, 1]
+        n = self.np_tensor[ops.convert_to_numpy(index_tensor), 1]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        index_tensor = ops.convert_to_tensor(np.array(-1, dtype=np.int32))
+        t = self.tensor[-2, index_tensor]
+        n = self.np_tensor[-2, ops.convert_to_numpy(index_tensor)]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        # Negative indexing
+        t = self.tensor[-1]
+        n = self.np_tensor[-1]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[1, -1, -2]
+        n = self.np_tensor[1, -1, -2]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        # Slicing with step
+        t = self.tensor[::2]
+        n = self.np_tensor[::2]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        # Mixed slices and integers
+        t = self.tensor[1, :, 1:4]
+        n = self.np_tensor[1, :, 1:4]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
+        t = self.tensor[:, 1:2, 3]
+        n = self.np_tensor[:, 1:2, 3]
+        self.assertEqual(t.shape, n.shape)
+        self.assertTrue(np.array_equal(ops.convert_to_numpy(t), n))
+
     def test_map(self):
         def f(x):
             return x**2