[Operator] Add Abs and And operator (#10)

python/hidet/graph/frontend/onnx.py

@@ -931,6 +931,18 @@ def run(self, inputs: List[Tensor]) -> List[Tensor]:
         return [ops.prelu(inputs[0], inputs[1])]
 
 
+@register_onnx_operator
+class OnnxAbs(OnnxOperator):
+    def run(self, inputs: List[Tensor]) -> List[Tensor]:
+        return [ops.abs(inputs[0])]
+
+
+@register_onnx_operator
+class OnnxAnd(OnnxOperator):
+    def run(self, inputs: List[Tensor]) -> List[Tensor]:
+        return [ops.cond_and(inputs[0], inputs[1])]
+
+
 def dispatch(node, op_sets: List[int]) -> OnnxOperator:
     op_type = node.op_type
     if op_type not in dispatch_table:

python/hidet/graph/ops/__init__.py

@@ -9,9 +9,9 @@
 from .definitions.activation import relu, leaky_relu, sigmoid, clip, relu6, prelu
 from .definitions.norm import batch_norm_infer, instance_norm, layer_norm
 from .definitions.image import resize2d
-from .definitions.arithmetic import add, sub, multiply, divide, neg, sqrt, rsqrt, sin, cos, pow, erf, tanh, where
+from .definitions.arithmetic import add, sub, multiply, divide, neg, sqrt, rsqrt, sin, cos, pow, erf, tanh, where, abs
 from .definitions.arithmetic import square, max, min, reciprocal, exp, log
-from .definitions.compare import equal, less_than, greater_than, less_or_equal, greater_or_equal, cond_not
+from .definitions.compare import equal, less_than, greater_than, less_or_equal, greater_or_equal, cond_not, cond_and
 from .definitions.reduce import reduce_mean, reduce_sum, reduce_var, reduce_min, reduce_max, argmin, argmax
 from .definitions.cumulative import cumsum
 from .definitions.transform import squeeze, unsqueeze, flatten, concat, cast, take, rearrange, strided_slice, reshape

python/hidet/graph/ops/definitions/__init__.py

@@ -11,8 +11,8 @@
 from .activation import relu, sigmoid, relu6, clip, prelu
 from .norm import batch_norm_infer, instance_norm
 from .image import resize2d
-from .arithmetic import add, sub, multiply, divide, neg, sqrt, rsqrt, where, max, min, reciprocal, exp, log
-from .compare import equal, less_than, greater_than, less_or_equal, greater_or_equal, cond_not
+from .arithmetic import add, sub, multiply, divide, neg, sqrt, rsqrt, where, max, min, reciprocal, exp, log, abs
+from .compare import equal, less_than, greater_than, less_or_equal, greater_or_equal, cond_not, cond_and
 from .reduce import reduce_mean, reduce_min, reduce_max, reduce_sum, reduce_var, argmin, argmax
 from .transform import squeeze, unsqueeze, flatten, concat, cast, take, rearrange, strided_slice, split, pad, conv_pad
 from .cumulative import cumsum

python/hidet/graph/ops/definitions/arithmetic.py

@@ -3,7 +3,7 @@
 
 from hidet.ir import primitives
 from hidet.ir import expr
-from hidet.ir.expr import const_like
+from hidet.ir.expr import const_like, if_then_else
 from hidet.utils import prod
 from hidet.graph.tensor import convert
 from .utils import Task, Operator, Tensor, TensorNode, InverseMap, compute, input_like
@@ -272,6 +272,11 @@ def scalar_min(args: List[expr.Expr]):
         )
 
 
+class AbsOp(UnaryElementwiseOp):
+    def __init__(self, x: Tensor):
+        super().__init__(x, op=lambda a: if_then_else(a >= 0, a, -a), name='abs')
+
+
 PythonScalar = Union[float, int]
 
 
@@ -406,3 +411,7 @@ def max(a: Tensor, b: Tensor, *others: Tensor) -> Tensor:
 def min(a: Tensor, b: Tensor, *others: Tensor) -> Tensor:
     args = [a, b] + list(others)
     return MinOp(args).get_output(0)
+
+
+def abs(x: Tensor) -> Tensor:
+    return AbsOp(x).get_output(0)

python/hidet/graph/ops/definitions/compare.py

@@ -33,6 +33,11 @@ def __init__(self, x: Tensor, y: Tensor):
         super().__init__(x, y, lambda a, b: a >= b, name='ge')
 
 
+class AndOp(BinaryElementwiseOp):
+    def __init__(self, x: Tensor, y: Tensor):
+        super().__init__(x, y, lambda a, b: expr.And(a, b), name='and')
+
+
 def cond_not(x: Tensor) -> Tensor:
     return NotOp(x).get_output(0)
 
@@ -57,3 +62,7 @@ def less_or_equal(x: Tensor, y: Tensor) -> Tensor:
 
 def greater_or_equal(x: Tensor, y: Tensor) -> Tensor:
     return GreaterOrEqual(x, y).get_output(0)
+
+
+def cond_and(x: Tensor, y: Tensor) -> Tensor:
+    return AndOp(x, y).get_output(0)

tests/graph/operators/test_arithmatic.py → tests/graph/operators/test_arithmetic.py

@@ -79,5 +79,10 @@ def test_neg(shape):
     check_unary(shape, np.float32, np.negative, ops.neg)
 
 
+@pytest.mark.parametrize("shape", unary_op_shapes)
+def test_abs(shape):
+    check_unary(shape, np.float32, np.absolute, ops.abs)
+
+
 if __name__ == '__main__':
     pytest.main([__file__])

tests/graph/operators/test_compare.py

@@ -0,0 +1,17 @@
+import numpy as np
+import pytest
+
+from hidet import ops
+from hidet.testing import check_unary, check_binary
+
+
+@pytest.mark.parametrize("a_shape, b_shape, dtype", [[[33, 44], [44], "bool"]])
+def test_and(a_shape, b_shape, dtype):
+    # without broadcast
+    check_binary(a_shape, a_shape, lambda a, b: np.logical_and(a, b), lambda a, b: ops.cond_and(a, b), dtype=dtype)
+    # with broadcast
+    check_binary(a_shape, b_shape, lambda a, b: np.logical_and(a, b), lambda a, b: ops.cond_and(a, b), dtype=dtype)
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])