Set TensorProxy shape's history to None if tensor's history is unavailable (#2755)

Co-authored-by: Masato Shinokawa <[email protected]>

Author: shino16

thunder/core/proxies.py

@@ -2011,12 +2011,21 @@ def tensorproxy(t: torch.Tensor, /, *, name: None | str, history: None | tuple =
     _thunder_fsdp_padding_size = getattr(t, "_thunder_fsdp_padding_size", None)
     # For parameters, shapes should be static.
     if using_symbolic_values() and not isinstance(t, torch.nn.Parameter):
-        shape_attr = ProvenanceRecord(PseudoInst.LOAD_ATTR, inputs=[copy.copy(history), wrap_const("shape").provenance])
+        if history is not None:
+            shape_pr = ProvenanceRecord(
+                PseudoInst.LOAD_ATTR, inputs=[copy.copy(history), wrap_const("shape").provenance]
+            )
+            dim_pr = lambda idx: ProvenanceRecord(
+                PseudoInst.BINARY_SUBSCR, inputs=[shape_pr, wrap_const(idx).provenance]
+            )
+        else:
+            dim_pr = lambda idx: None
+
         shape = tuple(
             IntegerProxy(
                 None,
                 s,
-                history=ProvenanceRecord(PseudoInst.BINARY_SUBSCR, inputs=[shape_attr, wrap_const(idx).provenance]),
+                history=dim_pr(idx),
                 constraint=CONSTRAINT.CONSTRAINABLE,
             )
             for idx, s in enumerate(t.shape)

thunder/tests/distributed/test_dtensor.py

@@ -13,7 +13,7 @@
 import thunder
 
 from thunder.tests.distributed.helper import DistributedParallelTestCase
-from torch.distributed._tensor import DeviceMesh, distribute_tensor
+from torch.distributed.tensor import DTensor, DeviceMesh, distribute_tensor
 from torch.distributed.tensor.placement_types import Shard, Replicate
 from torch.testing._internal.distributed._tensor.common_dtensor import DTensorConverter
 from torch.distributed.tensor.parallel import (
@@ -462,6 +462,33 @@ def test_dtensor_opinfo(self, op: OpInfo, executor):
 
         assert tested_sample_count > 0, f"test_dtensor_opinfo:No samples tested for {op.name} with {executor} executor"
 
+    def test_dtensor_from_local_symbolic_values(self):
+        num_devices = self.world_size
+        mesh = DeviceMesh("cuda", list(range(num_devices)))
+
+        dim_size = 8
+        local_tensor = torch.randn(dim_size, dim_size, device="cuda")
+
+        def fn(x):
+            return DTensor.from_local(x, mesh, [Shard(0)])
+
+        tjit = thunder.jit(fn, cache="symbolic values")
+
+        actual = tjit(local_tensor)
+        expected = DTensor.from_local(local_tensor, mesh, [Shard(0)])
+
+        torch.testing.assert_close(actual, expected)
+        assert thunder.cache_misses(tjit) == 1
+        assert thunder.cache_hits(tjit) == 0
+
+        dim_size = 16
+        local_tensor = torch.randn(dim_size, dim_size, device="cuda")
+        actual = tjit(local_tensor)
+        expected = DTensor.from_local(local_tensor, mesh, [Shard(0)])
+        torch.testing.assert_close(actual, expected)
+        assert thunder.cache_misses(tjit) == 1
+        assert thunder.cache_hits(tjit) == 1
+
 
 common_utils.instantiate_parametrized_tests(DTensorTest)