Fix PruneEmptyTensorsPass crash on data-dependent shapes (pytorch#20809)

### 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

Author: psiddh

exir/passes/prune_empty_tensors_pass.py

@@ -11,6 +11,7 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 from torch.fx import GraphModule, Node
+from torch.fx.experimental.symbolic_shapes import guard_or_true
 
 # This is a list of ops that are No Ops if used with an empty tensor.
 # Which means that if we remove the empty tensor as input to this op,
@@ -35,7 +36,7 @@ def remove_empty_tensors_from_cat(
         pruned_concat_list = []
         for input_arg in concat_list:
             input_arg_tensor = input_arg.meta["val"]
-            if input_arg_tensor.numel() != 0:
+            if guard_or_true(input_arg_tensor.numel() != 0):
                 pruned_concat_list.append(input_arg)
 
         cat_node.args = (pruned_concat_list,) + cat_node.args[1:]

exir/tests/test_prune_empty_tensors_pass.py

@@ -12,6 +12,7 @@
 from executorch.exir.capture._config import ExecutorchBackendConfig
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.passes import MemoryPlanningPass
+from executorch.exir.passes.prune_empty_tensors_pass import PruneEmptyTensorsPass
 
 
 class TestCat(nn.Module):
@@ -25,6 +26,16 @@ def get_example_inputs(self):
         return (torch.rand(5, 6), torch.rand(5, 6), torch.rand(5, 6))
 
 
+class UnbackedCatModel(nn.Module):
+    __test__ = False
+
+    def forward(self, x, routing_weights):
+        expert_mask = routing_weights > 0.5
+        indices = torch.nonzero(expert_mask)
+        selected = x[indices[:, 0]]
+        return torch.cat([x[:0], selected], dim=0)
+
+
 class TestPruneEmptyTensors(unittest.TestCase):
     def test_empty_tensor_removed_from_cat(self) -> None:
         model = TestCat()
@@ -77,3 +88,25 @@ def test_cat_removed_all_empty(self) -> None:
         reference = model(*example_inputs)
 
         self.assertTrue(torch.allclose(actual, reference))
+
+    def test_unbacked_symint_numel_does_not_crash(self) -> None:
+        """PruneEmptyTensorsPass must not crash on tensors whose numel() is an
+        unbacked SymInt (e.g. from torch.nonzero in MoE expert routing).
+        The pass should conservatively keep such tensors."""
+        model = UnbackedCatModel()
+        model.eval()
+        ep = torch.export.export(
+            model,
+            (torch.randn(4, 8), torch.randn(4, 2)),
+            strict=False,
+        )
+        result = PruneEmptyTensorsPass()(ep.graph_module)
+        found_cat = False
+        for node in result.graph_module.graph.nodes:
+            if node.target == torch.ops.aten.cat.default:
+                found_cat = True
+                cat_inputs = node.args[0]
+                self.assertEqual(len(cat_inputs), 1)
+                val = cat_inputs[0].meta["val"]
+                self.assertGreater(len(val.shape), 0)
+        self.assertTrue(found_cat)