Repository: pytorch
Issue: `CapabilityBasedPartitioner` returns invalid partitions.
Labels: triaged,module: xla,oncall: fx

### 🐛 Describe the bug

Running `Super_SloMo` from torchbench with XLA fails with:

```python
  File "torch/_dynamo/backends/torchxla.py", line 49, in fwd
    compiled_graph = bridge.extract_compiled_graph(model, args)
  File "xla/torch_xla/core/dynamo_bridge.py", line 523, in extract_compiled_graph
    partitioned_graph = partitioner.fuse_partitions(partitions)
  File "torch/fx/passes/infra/partitioner.py", line 217, in fuse_partitions
    return fuse_by_partitions(self.graph_module, [list(partition.nodes) for partition in partitions])
  File "torch/fx/passes/utils/fuser_utils.py", line 224, in fuse_by_partitions
    sub_gm, orig_inputs, orig_outputs = fuse_as_graphmodule(gm, sorted_nodes, submodule_name)
  File "torch/fx/passes/utils/fuser_utils.py", line 123, in fuse_as_graphmodule
    assert validate_partition(nodes), "Invalid partition, found dependency cycles"
AssertionError: Invalid partition, found dependency cycles
```

By applying the patch below, we can see that the list of partitions returned by `proposed_partitions` is already invalid:

```python
diff --git a/torch/fx/passes/infra/partitioner.py b/torch/fx/passes/infra/partitioner.py
index 7693f528af5..45080905cea 100644
--- a/torch/fx/passes/infra/partitioner.py
+++ b/torch/fx/passes/infra/partitioner.py
@@ -1,6 +1,6 @@
 from typing import Dict, List, Set, Iterable, Sequence, Optional, Deque

-from torch.fx.passes.utils.fuser_utils import fuse_by_partitions
+from torch.fx.passes.utils.fuser_utils import fuse_by_partitions, validate_partition

 from torch.fx.graph_module import GraphModule
 from torch.fx.node import Node, _get_qualified_name
@@ -208,6 +208,7 @@ class CapabilityBasedPartitioner:
         logger.debug("Partitions proposed:")
         for id, partition in partitions_by_id.items():
             logger.debug("partition #%s: %s", id, [node.name for node in partition.nodes])
+            assert validate_partition(list(partition.nodes)), "BAD"

         return list(partitions_by_id.values())
```

Command used: 
```
PJRT_DEVICE=CPU python benchmarks/dynamo/torchbench.py --randomize-input --performance --trace-on-xla --backend=openxla_eval --inference --only Super_SloMo
```

### Versions

- PyTorch version: 2bf3ca1be759460cf9fbf011d96d3246001361e9 (Oct 4)
- PyTorch/XLA version: c9a132484fb89bfdc9c602ada7bd8a3cec0db1aa (Oct 3)
- PyTorch/benchmark version: bb5275f819d479baacd24fe96000474355586096 (Oct 4)
    - Needs a patch in order to work with XLA

```python
diff --git a/torchbenchmark/models/Super_SloMo/slomo_model.py b/torchbenchmark/models/Super_SloMo/slomo_model.py
index a06e51eb..60160c18 100644
--- a/torchbenchmark/models/Super_SloMo/slomo_model.py
+++ b/torchbenchmark/models/Super_SloMo/slomo_model.py
@@ -249,8 +249,10 @@ class backWarp(nn.Module):

         # Use torch.meshgrid instead of np.meshgrid to imrpove performance
         # https://github.com/avinashpaliwal/Super-SloMo/pull/111
-        self.gridX, self.gridY = torch.meshgrid(torch.arange(W, requires_grad=False, device=device),
-                                                torch.arange(H, requires_grad=False, device=device), indexing='xy')
+        gridX, gridY = torch.meshgrid(torch.arange(W, requires_grad=False, device=device),
+                                      torch.arange(H, requires_grad=False, device=device), indexing='xy')
+        self.register_buffer("gridX", gridX)
+        self.register_buffer("gridY", gridY)

     def forward(self, img, flow):
         """

```

cc @bdhirsh @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv @msaroufim @wconstab @anijain2305 @zou3519