Add elu AIT converter (facebookincubator#953)

Summary:
Pull Request resolved: facebookincubator#953

ATT

Reviewed By: frank-wei, zoranzhao

Differential Revision: D50384768

fbshipit-source-id: c8a9787fb9313e3fcf712212fb27e2d9ad970a63

Author: aakhundov

fx2ait/fx2ait/converters/ait_converters.py

@@ -488,6 +488,31 @@ def acc_ops_relu(
     return elementwise(FuncEnum.RELU)(input_val)
 
 
+@ait_converter(acc_ops.elu)
+def acc_ops_elu(
+    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}")
+
+    inputs = [input_val]
+    if "alpha" in kwargs:
+        if not isinstance(kwargs["alpha"], (int, float)):
+            raise RuntimeError(
+                f"When specified, alpha in {name} must be a scalar: {input_val}"
+            )
+        input_alpha = AITTensor(
+            shape=[],
+            dtype=input_val._attrs["dtype"],
+            name="alpha",
+            value=kwargs["alpha"],
+        )
+        inputs.append(input_alpha)
+
+    return elementwise(FuncEnum.ELU)(*inputs)
+
+
 @ait_converter(acc_ops.leaky_relu)
 def acc_ops_leaky_relu(
     target: Target, args: Tuple[Argument, ...], kwargs: Dict[str, Argument], name: str

fx2ait/fx2ait/test/converters/test_ait_elu.py

@@ -0,0 +1,44 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+import torch
+from fx2ait.acc_tracer import acc_ops
+from fx2ait.tools.common_fx2ait import AITTestCase
+
+
+class TestEluConverter(AITTestCase):
+    def test_elu(self):
+        class TestModule(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return torch.nn.functional.elu(x)
+
+        model = TestModule().cuda().half()
+        inputs = [
+            torch.randn(2, 3).half().cuda(),
+        ]
+
+        self.run_test(model, inputs, expected_ops={acc_ops.elu})
+
+    def test_elu_with_alpha(self):
+        class TestModule(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return torch.nn.functional.elu(x, alpha=3.14)
+
+        model = TestModule().cuda().half()
+        inputs = [
+            torch.randn(2, 3).half().cuda(),
+        ]
+
+        self.run_test(model, inputs, expected_ops={acc_ops.elu})