diff --git a/backends/cadence/aot/tests/test_replace_ops_passes.py b/backends/cadence/aot/tests/test_replace_ops_passes.py index 6d1e34060ff..469214ed008 100644 --- a/backends/cadence/aot/tests/test_replace_ops_passes.py +++ b/backends/cadence/aot/tests/test_replace_ops_passes.py @@ -326,6 +326,33 @@ def test_sub_replace_scalar_with_tensor_arg( 0, ) + @expand( + [ + ["add", exir_ops.edge.aten.add.Scalar, exir_ops.edge.aten.add.Tensor], + ["sub", exir_ops.edge.aten.sub.Scalar, exir_ops.edge.aten.sub.Tensor], + ] + ) + @torch.no_grad() + def test_replace_scalar_with_tensor_arg_positional_alpha( + self, + _, + scalar_op: torch._ops.OpOverload, + tensor_op: torch._ops.OpOverload, + ) -> None: + x = torch.randn(7, 5) + original_gm = single_op_builder( + placeholders=(x,), + op=scalar_op, + args=(x, 1.5, 2), + ) + gm_before = copy.deepcopy(original_gm) + p = ReplaceScalarWithTensorArgPass() + graph_after_passes = cast(PassResult, p(original_gm)).graph_module + + validate(gm_before, graph_after_passes, [x], "ReplaceScalarWithTensorArgPass") + self.assertEqual(count_node(graph_after_passes, tensor_op), 1) + self.assertEqual(count_node(graph_after_passes, scalar_op), 0) + @expand( [ ["3d", (7, 5, 6), 1.23], diff --git a/backends/transforms/replace_scalar_with_tensor.py b/backends/transforms/replace_scalar_with_tensor.py index ed45998dc56..3ff83e0e3bf 100644 --- a/backends/transforms/replace_scalar_with_tensor.py +++ b/backends/transforms/replace_scalar_with_tensor.py @@ -12,6 +12,7 @@ from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.pass_base import ExportPass from torch._ops import OpOverload +from torch.fx.operator_schemas import normalize_function class ReplaceScalarWithTensorArgPass(ExportPass): @@ -71,8 +72,6 @@ def get_replacement(self, op, args, kwargs, meta): }, meta=meta, ), - # Other args. - *args[2:], ), kwargs=kwargs, meta=meta, @@ -96,4 +95,10 @@ def call_operator(self, op, args, kwargs, meta): ): return super().call_operator(op, args, kwargs, meta) - return self.get_replacement(op, args, kwargs, meta) + # alpha is positional in add/sub.Scalar but keyword-only in .Tensor. + bound = normalize_function( + op, args, kwargs, normalize_to_only_use_kwargs=True + ).kwargs + return self.get_replacement( + op, (bound["input"], bound["other"]), {"alpha": bound["alpha"]}, meta + )