Skip to content

Commit a2ae4b3

Browse files
Qualcomm AI Engine Direct - Preserve partition metadata on single-output delegate getitem nodes (pytorch#20805)
### Summary The single-output path in `create_submodule_from_nodes` copied only "val" onto the manually-built getitem, dropping other partition meta (e.g. QNN's `q_tensor_io` tag) that the multi-output path keeps via `propagate_meta`, leaving backend IO-tagging passes with wrong delegate `dtype`/`spec`. Fix by propagating full meta then dropping `nn_module_stack`/`source_fn_stack`, matching the multi-output path.
1 parent 7a6f630 commit a2ae4b3

1 file changed

Lines changed: 3 additions & 2 deletions

File tree

‎exir/lowered_backend_module.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -884,10 +884,11 @@ def create_submodule_from_nodes(
884884
# all uses with a getitem call to the 0th index of the result
885885
with gm.graph.inserting_after(submodule_node):
886886
proxy_out = torch.fx.Proxy(submodule_node)[0].node # type: ignore[index]
887-
submodule_node.replace_all_uses_with(proxy_out)
888-
proxy_out.meta["val"] = submodule_node.meta["val"]
887+
submodule_node.replace_all_uses_with(proxy_out, propagate_meta=True)
889888
# Reset the args since it was overwritten in the previous line
890889
proxy_out.args = (submodule_node, 0)
890+
proxy_out.meta.pop("nn_module_stack", None)
891+
proxy_out.meta.pop("source_fn_stack", None)
891892
else:
892893
# fuse_as_graphmodule will automatically propagate the metadata of the
893894
# partition's last node to the getitem nodes that appear after the

0 commit comments

Comments
 (0)