[Unity] [Relax] [ONNX frontend] [op] Add support for Trilu operator

python/tvm/relax/frontend/onnx/onnx_frontend.py

@@ -499,6 +499,23 @@ def _impl_v13(cls, bb, inputs, attr, params):
         return relax.op.sqrt(inputs[0])
 
 
+class Trilu(OnnxOpConverter):
+    """Given a 2-D matrix or batches of 2-D matrices, returns the upper or
+    lower triangular part of the tensor(s)
+    """
+
+    @classmethod
+    def _impl_v14(cls, bb, inputs, attr, params):
+        upper = attr.get("upper", True)
+        x = inputs[0]
+        k = inputs[1] if len(inputs) > 1 else 0
+
+        if upper:
+            return relax.op.triu(x, k)
+        else:
+            return relax.op.tril(x, k)
+
+
 class Relu(OnnxOpConverter):
     """Converts an onnx Relu node into an equivalent Relax expression."""
 
@@ -1712,6 +1729,7 @@ def _get_convert_map():
         "Shape": Shape,
         "Tanh": Tanh,
         "Sqrt": Sqrt,
+        "Trilu": Trilu,
         "Relu": Relu,
         "Conv": Conv,
         "Pow": Pow,

tests/python/relax/test_frontend_onnx.py

@@ -528,6 +528,14 @@ def test_relu():
     verify_unary("Relu", [32, 32])
 
 
+def test_tril():
+    verify_unary("Trilu", [3, 5, 5], attrs={"upper": False})
+
+
+def test_triu():
+    verify_unary("Trilu", [3, 5, 5], attrs={"upper": True})
+
+
 def test_conv():
     def _verify_conv(input_shape, weight_shape, output_shape):
         bias_shape = [output_shape[1]]