From 9e147d7c2334c2dcbd74217b137ab5cdd244c9cb Mon Sep 17 00:00:00 2001 From: Ali Afzal Date: Wed, 28 Oct 2026 13:57:02 -0700 Subject: [PATCH 1/2] Move positional alpha to kwargs in ReplaceScalarWithTensorArgPass (#24176) Summary: `add.Scalar(Tensor self, Scalar other, Scalar alpha=1)` takes `alpha` positionally, but `add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1)` takes it as keyword-only, and the same holds for `sub`. When `alpha` was the third positional argument, `get_replacement` passed it through `*args[2:]`, and building the `.Tensor` node raised `RuntimeError: aten::add() takes 2 positional argument(s) but 3 was/were given`. The 3-argument path now binds the call by name with `torch.fx.operator_schemas.normalize_function` (the part of `get_arg` that doesn't need a `Node`; `call_operator` gets no node), passes `self` (`input` after normalization) and `other` positionally, and `alpha` as a kwarg. `get_replacement` always gets 2 args and no longer forwards extra ones. The Arm `ReplaceScalarWithTensorByProfilePass` subclass routes through the same base `call_operator`, so it gets the fix too for edge `add/sub.Scalar`. `torch.ops.aten.add/sub.Scalar` with a positional `alpha` still pass through unreplaced, as before. The changed path previously always raised, so no graph that worked before is affected. Authored with Claude Code. Reviewed By: DrJessop Differential Revision: D127746753 --- .../aot/tests/test_replace_ops_passes.py | 27 +++++++++++++++++++ .../transforms/replace_scalar_with_tensor.py | 11 +++++--- 2 files changed, 35 insertions(+), 3 deletions(-) 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 + ) From 320e66733ad0026c62b4cd1d0c91454ef578e700 Mon Sep 17 00:00:00 2001 From: Ali Afzal Date: Wed, 28 Oct 2026 13:57:02 -0700 Subject: [PATCH 2/2] Skip ReplacePowWithMulPass when pow promotes the input dtype (#24177) Summary: `pow.Tensor_Scalar(int32, 2.0)` returns float32, but the `mul.Tensor(x, x)` chain that `ReplacePowWithMulPass` emits stays int32. The pass copied the float32 `meta` onto the int32 `mul` nodes, so the lowered program's output dtype no longer matched the eager reference. The pass now gives up when the input dtype differs from the `pow` output dtype. Authored with Claude Code. Reviewed By: DrJessop Differential Revision: D127746754 --- backends/cadence/aot/replace_ops.py | 5 ++ .../aot/tests/test_replace_ops_passes.py | 47 +++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/backends/cadence/aot/replace_ops.py b/backends/cadence/aot/replace_ops.py index f68b05857d2..bc04f749491 100644 --- a/backends/cadence/aot/replace_ops.py +++ b/backends/cadence/aot/replace_ops.py @@ -2280,6 +2280,11 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: x = node.args[0] assert isinstance(x, torch.fx.Node) + # mul keeps the input dtype, so it can't express a promoting pow + # (e.g. int32 ** 2.0 -> float32). + if x.meta["val"].dtype != node.meta["val"].dtype: + return False + graph = node.graph result_node = x diff --git a/backends/cadence/aot/tests/test_replace_ops_passes.py b/backends/cadence/aot/tests/test_replace_ops_passes.py index 469214ed008..0b6c07bddee 100644 --- a/backends/cadence/aot/tests/test_replace_ops_passes.py +++ b/backends/cadence/aot/tests/test_replace_ops_passes.py @@ -2375,6 +2375,53 @@ def test_replace_pow_with_mul_not_applied(self, exponent: float) -> None: 0, ) + @expand([[2.0], [3.0]]) + def test_replace_pow_with_mul_not_applied_on_type_promotion( + self, exponent: float + ) -> None: + # int32 ** float promotes to float32, but int32 * int32 stays int32. + x = torch.randint(0, 8, (2, 1, 64), dtype=torch.int32) + original_gm = single_op_builder( + placeholders=(x,), + op=exir_ops.edge.aten.pow.Tensor_Scalar, + args=(x, exponent), + ) + p = ReplacePowWithMulPass() + result = cast(PassResult, p(original_gm)) + + self.assertFalse(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.pow.Tensor_Scalar), + 1, + ) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.mul.Tensor), 0 + ) + + @expand([[2], [3]]) + def test_replace_pow_with_mul_int_input_int_exponent(self, exponent: int) -> None: + # int32 ** int stays int32, so the dtype guard still allows the rewrite. + x = torch.randint(0, 8, (2, 1, 64), dtype=torch.int32) + original_gm = single_op_builder( + placeholders=(x,), + op=exir_ops.edge.aten.pow.Tensor_Scalar, + args=(x, exponent), + ) + gm_before = copy.deepcopy(original_gm) + p = ReplacePowWithMulPass() + result = cast(PassResult, p(original_gm)) + + self.assertTrue(result.modified) + validate(gm_before, result.graph_module, [x], "ReplacePowWithMulPass") + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.pow.Tensor_Scalar), + 0, + ) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.mul.Tensor), + exponent - 1, + ) + class TestReplaceIm2rowWithViewPass(unittest.TestCase): def test_no_replacement_for_conv(self) -> None: