[Graph][Dynamo Backend] Lshift/Rshift/Mod (#371)

Add operators lshift, rshift, mod
hidet-org · Oct 29, 2023 · d9a142a · d9a142a
1 parent 0d4ce21
commit d9a142a
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/python/hidet/graph/frontend/torch/register_functions.py b/python/hidet/graph/frontend/torch/register_functions.py
@@ -1024,6 +1024,21 @@ def ne(a: Union[Tensor, Expr, Number], b: Union[Tensor, Expr, Number]) -> Tensor
     return a != b
 
 
+@register_function(operator.mod)
+def mod(a: Union[Tensor, Expr, Number], b: Union[Tensor, Expr, Number]) -> Tensor:
+    return a % b
+
+
+@register_function(operator.lshift)
+def lshift(a: Union[Tensor, Expr, Number], b: Union[Tensor, Expr, Number]) -> Tensor:
+    return a << b
+
+
+@register_function(operator.rshift)
+def rshift(a: Union[Tensor, Expr, Number], b: Union[Tensor, Expr, Number]) -> Tensor:
+    return a >> b
+
+
 @register_function(torch.rsqrt)
 def rsqrt(x: Tensor, *, out: Optional[Tensor] = None) -> Tensor:
     if out is not None: