From a3545fe0b342c5528416153394e32f72e3a1f295 Mon Sep 17 00:00:00 2001 From: Andrew Grebenisan Date: Tue, 9 Dec 2025 12:15:48 -0800 Subject: [PATCH] Fix memory planning error Summary: Added back some missing code during previous refactor the pass. Without it, was hitting a memory planning failure. Differential Revision: D88770975 --- backends/cadence/aot/replace_ops.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/backends/cadence/aot/replace_ops.py b/backends/cadence/aot/replace_ops.py index 0b88e687224..fa8e8ca202a 100644 --- a/backends/cadence/aot/replace_ops.py +++ b/backends/cadence/aot/replace_ops.py @@ -370,15 +370,23 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: graph = node.graph + fit_bias = is_node_with_op(bias, "get_attr") or beta == 1.0 + fit_mat2 = is_node_with_op(mat2, "get_attr") + # Handle transpose: if mat2 is a transpose op, extract the original tensor transposed_mat2 = False if ( - mat2.op == "call_function" + not fit_mat2 + and mat2.op == "call_function" and mat2.target == exir_ops.edge.aten.transpose_copy.int ): # mat2 is already transposed, so we use the input to the transpose mat2 = cast(torch.fx.Node, mat2.args[0]) transposed_mat2 = True + fit_mat2 = is_node_with_op(mat2, "get_attr") or alpha == 1.0 + + if not (fit_bias and fit_mat2): + return False # Multiply bias by beta if needed if beta != 1.0: