Repository: pytorch
Issue: "unhashable type: non-nested SymInt" error when using DTensor and Compiled Autograd together
Labels: triaged,module: fsdp,oncall: pt2,module: dynamic shapes,module: dtensor,module: compiled autograd

### 🐛 Describe the bug

Repro:

If I make the following change to enable Compiled Autograd for compile+DTensor E2E test `test_tp_compile_fullgraph`:

```
diff --git a/test/distributed/_tensor/test_dtensor_compile.py b/test/distributed/_tensor/test_dtensor_compile.py
index 325d18be79f..f142314c660 100644
--- a/test/distributed/_tensor/test_dtensor_compile.py
+++ b/test/distributed/_tensor/test_dtensor_compile.py
@@ -661,9 +661,13 @@ class TestDTensorCompileE2E(DTensorTestBase):
         inp = torch.rand(20, 10, device=self.device_type)
         out = model(inp)
         cnt = torch._dynamo.testing.CompileCounterWithBackend("aot_eager")
-        compiled_mod = torch.compile(model, backend=cnt, fullgraph=True)
-        compiled_out = compiled_mod(inp)
-        compiled_out.sum().backward()
+        compiled_mod = torch.compile(model, backend="aot_eager", fullgraph=True)
+        def compiler_fn(gm):
+            return torch.compile(gm, backend="aot_eager", fullgraph=True)
+        ctx = torch._dynamo.compiled_autograd.enable(compiler_fn)
+        with ctx:
+            compiled_out = compiled_mod(inp)
+            compiled_out.sum().backward()
         self.assertEqual(compiled_out, out)
         self.assertEqual(cnt.frame_count, 1)
```
and run `pytest -rA test/distributed/_tensor/test_dtensor_compile.py::TestDTensorCompileE2E::test_tp_compile_fullgraph_is_seq_parallel_False`, I get this error:
```
Traceback (most recent call last):
  File "/data/users/willfeng/pytorch/torch/testing/_internal/common_distributed.py", line 653, in run_test
    getattr(self, test_name)()
  File "/data/users/willfeng/pytorch/torch/testing/_internal/common_distributed.py", line 539, in wrapper
    fn()
  File "/data/users/willfeng/pytorch/torch/testing/_internal/common_utils.py", line 2715, in wrapper
    method(*args, **kwargs)
  File "/data/users/willfeng/pytorch/torch/testing/_internal/common_utils.py", line 446, in instantiated_test
    test(self, **param_kwargs)
  File "/data/users/willfeng/pytorch/torch/testing/_internal/distributed/_tensor/common_dtensor.py", line 369, in wrapper
    func(self, *args, **kwargs)  # type: ignore[misc]
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/willfeng/pytorch/test/distributed/_tensor/test_dtensor_compile.py", line 670, in test_tp_compile_fullgraph
    compiled_out.sum().backward()
  File "/data/users/willfeng/pytorch/torch/_tensor.py", line 520, in backward
    torch.autograd.backward(
  File "/data/users/willfeng/pytorch/torch/autograd/__init__.py", line 284, in backward
    _engine_run_backward(
  File "/data/users/willfeng/pytorch/torch/autograd/graph.py", line 767, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/willfeng/pytorch/torch/_dynamo/compiled_autograd.py", line 88, in begin_capture
    self.bind_tensors_to_proxies(inputs, proxies)
  File "/data/users/willfeng/pytorch/torch/_dynamo/compiled_autograd.py", line 258, in bind_tensors_to_proxies
    track_tensor_tree(tensors, proxies, constant=None, tracer=self.fx_tracer)
  File "/data/users/willfeng/pytorch/torch/fx/experimental/proxy_tensor.py", line 267, in track_tensor_tree
    wrap_with_proxy(inner_res, proxy_res, constant)
  File "/data/users/willfeng/pytorch/torch/fx/experimental/proxy_tensor.py", line 238, in wrap_with_proxy
    wrap_with_proxy(ee, proxy[idx], get_constant(idx))
  File "/data/users/willfeng/pytorch/torch/fx/experimental/proxy_tensor.py", line 224, in wrap_with_proxy
    set_meta(proxy, e)
  File "/data/users/willfeng/pytorch/torch/fx/experimental/proxy_tensor.py", line 178, in set_meta
    proxy.node.meta['val'] = extract_val(val)
                             ^^^^^^^^^^^^^^^^
  File "/data/users/willfeng/pytorch/torch/fx/experimental/proxy_tensor.py", line 146, in extract_val
    return snapshot_fake(val)
           ^^^^^^^^^^^^^^^^^^
  File "/data/users/willfeng/pytorch/torch/fx/experimental/proxy_tensor.py", line 142, in snapshot_fake
    return val.detach()
           ^^^^^^^^^^^^
  File "/data/users/willfeng/pytorch/torch/_compile.py", line 30, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/willfeng/pytorch/torch/_dynamo/eval_frame.py", line 574, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/data/users/willfeng/pytorch/torch/distributed/_tensor/api.py", line 296, in __torch_dispatch__
    return DTensor._op_dispatcher.dispatch(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/willfeng/pytorch/torch/distributed/_tensor/_dispatch.py", line 117, in dispatch
    self.sharding_propagator.propagate(op_info)
  File "/data/users/willfeng/pytorch/torch/distributed/_tensor/_sharding_prop.py", line 182, in propagate
    output_sharding = self.propagate_op_sharding_non_cached(op_info.schema)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/willfeng/pytorch/torch/distributed/_tensor/_sharding_prop.py", line 196, in propagate_op_sharding_non_cached
    out_tensor_meta = self._propagate_tensor_meta(op_schema)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/willfeng/pytorch/torch/distributed/_tensor/_op_schema.py", line 333, in __hash__
    return hash((self.op, args_to_hash))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/willfeng/pytorch/torch/distributed/_tensor/placement_types.py", line 527, in __hash__
    self._hash = self._hash_impl()
                 ^^^^^^^^^^^^^^^^^
  File "/data/users/willfeng/pytorch/torch/distributed/_tensor/placement_types.py", line 510, in _hash_impl
    return hash(
           ^^^^^
  File "/data/users/willfeng/pytorch/torch/__init__.py", line 356, in __hash__
    raise TypeError("unhashable type: non-nested SymInt")
TypeError: unhashable type: non-nested SymInt
```

Not sure whether we must have this SymInt, or if we need a way to hash this SymInt.

### Versions

nightly

cc @zhaojuanmao @mrshenli @rohan-varma @awgu @fegin @kwen2501 @chauhang @ezyang @anijain2305 @penguinwu @wanchaol @XilunWu @d4l3k @msaroufim @bdhirsh @tianyu-l