Add acc_ops.exp support (#820)

Summary: Pull Request resolved: #820 Reviewed By: frank-wei Differential Revision: D47348835 fbshipit-source-id: 7e2a9645c174d8e7ae7f65fc2d450e1ae58b3c0f
facebookincubator · Jul 11, 2023 · 4667de0 · 4667de0
1 parent c016f6d
commit 4667de0
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/fx2ait/fx2ait/converters/ait_converters.py b/fx2ait/fx2ait/converters/ait_converters.py
@@ -382,6 +382,20 @@ def acc_ops_abs(
     return elementwise(FuncEnum.ABS)(input_val)
 
 
+@ait_converter(acc_ops.exp)
+def acc_ops_exp(
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> ConverterOutput:
+    input_val = kwargs["input"]
+    if not isinstance(input_val, AITTensor):
+        raise RuntimeError(f"Unexpected input for {name}: {input_val}")
+
+    return elementwise(FuncEnum.EXP)(input_val)
+
+
 @ait_converter(acc_ops.log)
 def acc_ops_log(
     target: Target,

diff --git a/fx2ait/fx2ait/test/converters/test_ait_unary_ops.py b/fx2ait/fx2ait/test/converters/test_ait_unary_ops.py
@@ -37,6 +37,7 @@
     (torch.sqrt, acc_ops.sqrt),
     (torch.clone, acc_ops.clone),
     (torch.neg, acc_ops.neg),
+    (torch.exp, acc_ops.exp),
 ]
 
 TestEnvToPrecision: Dict[TestEnv, Set[LowerPrecision]] = {