add trilu

Author: xadupre

_unittests/ut_reference/test_torch_onnx_evaluator.py

@@ -4,7 +4,7 @@
 import onnx.helper as oh
 import onnx.numpy_helper as onh
 import torch
-from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings
 from onnx_diagnostic.helpers.onnx_helper import from_array_extended
 from onnx_diagnostic.reference import ExtendedReferenceEvaluator, TorchOnnxEvaluator
 from onnx_diagnostic.reference.torch_evaluator import get_kernels
@@ -801,6 +801,56 @@ def test_op_constant_of_shape(self):
         onnx.checker.check_model(model)
         self._finalize_test(model, torch.tensor([4, 5], dtype=torch.int64))
 
+    def test_op_trilu(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("Trilu", ["X"], ["Z"])],
+                "dummy",
+                [oh.make_tensor_value_info("X", TFLOAT, ["a", "b"])],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        onnx.checker.check_model(model)
+        self._finalize_test(model, torch.rand((4, 4), dtype=torch.float32))
+
+    def test_op_trilu_1(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("Trilu", ["X"], ["Z"], upper=0)],
+                "dummy",
+                [oh.make_tensor_value_info("X", TFLOAT, ["a", "b"])],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        onnx.checker.check_model(model)
+        self._finalize_test(model, torch.rand((4, 4), dtype=torch.float32))
+
+    @ignore_warnings(DeprecationWarning)
+    def test_op_trilu_k(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("Trilu", ["X", "k"], ["Z"], upper=1)],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, ["a", "b"]),
+                    oh.make_tensor_value_info("k", TINT64, []),
+                ],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        onnx.checker.check_model(model)
+        self._finalize_test(
+            model,
+            torch.rand((6, 6), dtype=torch.float32),
+            torch.tensor([2], dtype=torch.int64),
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/reference/torch_evaluator.py

@@ -152,7 +152,7 @@ def _build_kernels(self, nodes: Sequence[onnx.NodeProto]):
             ), f"Missing kernel for node type {node.op_type!r} from domain {node.domain!r}"
             cls = kernels[key]
             if cls.device_dependent():
-                kernel = cls(node, opset, self.default_device)
+                kernel = cls(node, opset, self.default_device)  # type: ignore[call-arg]
             else:
                 kernel = cls(node, opset)
             self.kernels.append(kernel)

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -16,7 +16,7 @@
 )
 from .generator_ops import Range_11
 from .nn_ops import LayerNormalization_17, Softmax_13, Tanh_6
-from .other_ops import Cast_6, Concat_1, Transpose_1, Where_9
+from .other_ops import Cast_6, Concat_1, Transpose_1, Trilu_14, Where_9
 from .reduce_ops import ReduceMax_18, ReduceMean_18, ReduceMin_18, ReduceSum_18
 from .shape_ops import (
     ConstantOfShape_9,

onnx_diagnostic/reference/torch_ops/_op_run.py

@@ -61,7 +61,7 @@ class OpRun:
     @classmethod
     def device_dependent(cls) -> bool:
         """
-        Returns True if the kernel needs a device to be efficiently initiliazed.
+        Returns True if the kernel needs a device to be efficiently initialized.
         """
         return False
 

onnx_diagnostic/reference/torch_ops/other_ops.py

@@ -42,6 +42,20 @@ def run(self, data: OpRunValue) -> OpRunValue:
         return OpRunValue(torch.permute(data.tensor, self.perm))
 
 
+class Trilu_14(OpRun):
+    "Trilu"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        self.upper = self.get_attribute_int(node, "upper", 1)
+
+    def run(self, data: OpRunValue, k: Optional[OpRunValue] = None) -> OpRunValue:
+        diagonal = 0 if k is None else k.tensor.item()
+        if self.upper:
+            return OpRunValue(torch.triu(data.tensor, diagonal=diagonal))
+        return OpRunValue(torch.tril(data.tensor, diagonal=diagonal))
+
+
 class Where_9(OpRun):
     "Where"
 

onnx_diagnostic/reference/torch_ops/shape_ops.py

@@ -10,7 +10,7 @@ class ConstantOfShape_9(OpRun):
     @classmethod
     def device_dependent(cls) -> bool:
         """
-        Returns True if the kernel needs a device to be efficiently initiliazed.
+        Returns True if the kernel needs a device to be efficiently initialized.
         """
         return True