add gather

Author: xadupre

_unittests/ut_reference/test_torch_evaluator.py

@@ -240,6 +240,26 @@ def test_op_concat(self):
             torch.rand((4, 5, 2, 7), dtype=torch.float32),
         )
 
+    def test_op_gather(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("Gather", ["X", "Y"], ["Z"], axis=1)],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, ["a", "b", "c", "d"]),
+                    oh.make_tensor_value_info("Y", TINT64, ["f"]),
+                ],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b", "d"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        self._finalize_test(
+            model,
+            torch.rand((5, 4, 3, 2), dtype=torch.float32),
+            torch.tensor([0, 1, 3], dtype=torch.int64),
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -1,5 +1,5 @@
 from ._op_run import OpRun, OpRunValue
-from .access_ops import Slice_13
+from .access_ops import Gather_1, Slice_13
 from .binary_ops import Add_1, Div_1, MatMul_1, Mul_1, Sub_1
 from .other_ops import Cast_6, Concat_1, Transpose_1
 from .shape_ops import Reshape_5, Shape_15, Squeeze_13, Unsqueeze_13

onnx_diagnostic/reference/torch_ops/access_ops.py

@@ -1,7 +1,26 @@
 from typing import Optional
+import onnx
+import torch
 from . import OpRun, OpRunValue
 
 
+class Gather_1(OpRun):
+    "Gather"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        axis = self.get_attribute_int(node, "axis", 0)
+        assert isinstance(axis, int), f"Unexpected value for attribute axis={axis!r}"
+        self.axis = axis
+
+    def run(self, x, indices):
+        if indices.tensor.numel() == 0:
+            return torch.empty((0,), dtype=x.tensor.dtype, device=x.tensor.device)
+        ind = [slice(0, s) for s in x.shape]
+        ind[self.axis] = indices.tensor
+        return OpRunValue(x.tensor[*ind])
+
+
 class Slice_13(OpRun):
     "Slice"
 

onnx_diagnostic/reference/torch_ops/other_ops.py

@@ -6,6 +6,8 @@
 
 
 class Cast_6(OpRun):
+    "Cast"
+
     def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
         super().__init__(node, version)
         to = self.get_attribute_int(node, "to", 0)
@@ -17,6 +19,8 @@ def run(self, data: OpRunValue) -> OpRunValue:
 
 
 class Concat_1(OpRun):
+    "Concat"
+
     def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
         super().__init__(node, version)
         axis = self.get_attribute_int(node, "axis", 0)
@@ -28,6 +32,8 @@ def run(self, *data: OpRunValue) -> OpRunValue:
 
 
 class Transpose_1(OpRun):
+    "Transpose"
+
     def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
         super().__init__(node, version)
         self.perm = self.get_attribute_ints(node, "perm", None)