Fix PruneEmptyTensorsPass crash on data-dependent shapes - #20809
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20809
Note: Links to docs will display an error until the docs builds have been completed. ⏳ 1 Pending, 1 Unrelated FailureAs of commit 68e9ca2 with merge base f4b01a8 ( BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
|
|
This PR needs a
|
e06e55a to
3959dcd
Compare
There was a problem hiding this comment.
Pull request overview
This PR aims to prevent PruneEmptyTensorsPass from crashing when it encounters data-dependent (unbacked) symbolic shapes (e.g., shapes arising from nonzero/MoE routing), and adds end-to-end support for emitting/evaluating torch.sym_ite via an ExecuTorch prim op.
Changes:
- Update
PruneEmptyTensorsPassto handle symbolicnumel()comparisons without crashing on unbacked symbols. - Add
executorch_prim::sym_ite.Scalarregistration + runtime kernel, and register the corresponding EXIR prim op mapping. - Add C++ and Python tests covering
sym_iteregistration/runtime behavior and emit-to-ExecuTorch correctness.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| kernels/prim_ops/test/prim_ops_test.cpp | Adds registration and runtime-value tests for executorch_prim::sym_ite.Scalar. |
| kernels/prim_ops/register_prim_ops.cpp | Registers the executorch_prim::sym_ite.Scalar runtime kernel (selective-build gated). |
| exir/tests/test_prune_empty_tensors_pass.py | Adds a regression test for unbacked/data-dependent shapes in PruneEmptyTensorsPass. |
| exir/passes/prune_empty_tensors_pass.py | Changes empty-tensor pruning condition to avoid crashes on symbolic numel(). |
| exir/passes/executorch_prim_ops_registry.py | Introduces EXIR backend op schema/binding and mapping for torch.sym_ite. |
| exir/emit/test/test_emit.py | Adds an end-to-end emit/export test for torch.sym_ite with dynamic shapes. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
3959dcd to
6cb8730
Compare
6cb8730 to
75753c7
Compare
PruneEmptyTensorsPass crashes with GuardOnDataDependentSymNode when cat inputs have unbacked dynamic shapes (e.g. from torch.nonzero in MoE expert routing). The numel() check returns a SymBool that can't be evaluated as a Python bool. Use guard_or_true to conservatively keep tensors when emptiness can't be statically determined — this pass is purely an optimization, so keeping a potentially-empty tensor is always safe. This removes the need for HuggingFace's _patch_remove_empty_tensors_from_cat monkey-patch in their ExecuTorch exporter.
75753c7 to
68e9ca2
Compare
| if guard_or_true(input_arg_tensor.numel() != 0): | ||
| pruned_concat_list.append(input_arg) |
Summary
PruneEmptyTensorsPass crashes during export when cat inputs have data-dependent shapes (e.g. from conservatively keep the tensor. This is safe because the pass is purely an optimization; keeping an extra tensor in cat is a no-op at runtime.
This removes the need for HuggingFace's _patch_remove_empty_tensors_from_cat monkey-patch in their ExecuTorch exporter:(https://github.com/huggingface/transformers/blob/main/src/transformers/exporters/exporter_executorch.py#L456-L485).
Test plan
Existing tests pass: test_empty_tensor_removed_from_cat, test_cat_removed_all_empty. New test test_unbacked_symint_numel_does_not_crash exports a model using torch.nonzero (produces unbacked SymInt shapes) and verifies the pass completes without crashing
Verified guard_or_true returns True for undecidable expressions and correct concrete values for decidable ones
python -m pytest exir/tests/test_prune_empty_tensors_pass.py -v