[Fix][Relax][Frontend][Torch] Apply the scale and input_scale arguments of aten.elu

python/tvm/relax/frontend/torch/exported_program_translator.py

@@ -74,6 +74,28 @@ def _convert_pytorch_tensor_to_tvm(tensor_value: torch.Tensor) -> tvm.runtime.Te
 
     ########## Unary Ops ##########
 
+    def _elu(self, node: fx.Node) -> relax.Expr:
+        # aten.elu is elu(x, alpha, scale, input_scale); run_decompositions rewrites selu to it
+        scale = node.args[2] if len(node.args) > 2 else node.kwargs.get("scale", 1.0)
+        input_scale = node.args[3] if len(node.args) > 3 else node.kwargs.get("input_scale", 1.0)
+        if scale == 1 and input_scale == 1:
+            return super()._elu(node)
+
+        x = self.env[node.args[0]]
+        alpha = node.args[1] if len(node.args) > 1 else node.kwargs.get("alpha", 1.0)
+        dtype = x.ty.dtype
+        bb = self.block_builder
+        scaled_x = x
+        if input_scale != 1:
+            scaled_x = bb.emit(relax.op.multiply(x, relax.const(input_scale, dtype)))
+        # scale * (ReLU(x) - alpha * ReLU(1 - exp(input_scale * x)))
+        negative = relax.op.multiply(
+            relax.const(-alpha, dtype),
+            relax.op.nn.relu(relax.op.subtract(relax.const(1, dtype), relax.op.exp(scaled_x))),
+        )
+        out = bb.emit(relax.op.add(negative, relax.op.nn.relu(x)))
+        return out if scale == 1 else bb.emit(relax.op.multiply(out, relax.const(scale, dtype)))
+
     def _hardtanh(self, node: fx.Node) -> relax.Expr:
         args = self.retrieve_args(node)
         x = args[0]

tests/python/relax/test_frontend_from_exported_program.py

@@ -236,6 +236,21 @@ def main(input_1: R.Tensor((1, 3), dtype="int32")) -> R.Tuple(
     verify_model(SqrtIntModel(), example_args_int32, {}, expected_int32)
 
 
+def test_selu_and_elu_scale_arguments():
+    # run_decompositions rewrites selu to aten.elu(x, alpha, scale), so scale must be applied.
+    class Selu(Module):
+        def forward(self, x):
+            return torch.nn.functional.selu(x)
+
+    class EluScaled(Module):
+        def forward(self, x):
+            return torch.ops.aten.elu(x, 0.5, 2.0, 3.0)
+
+    example_args = (torch.tensor([[-2.0, -0.5, 0.0, 1.5]], dtype=torch.float32),)
+    verify_model_numerically(Selu(), example_args, rtol=1e-5, atol=1e-5)
+    verify_model_numerically(EluScaled(), example_args, rtol=1e-5, atol=1e-5)
+
+
 def test_extended_unary_ops():
     example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
 
@@ -706,7 +721,10 @@ def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
                 )
                 lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(input)
                 lv5: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(lv3, lv4)
-                gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv5,)
+                lv6: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(
+                    lv5, R.const(1.0507009873554805, "float32")
+                )
+                gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv6,)
                 R.output(gv)
             return gv