Symbolic values: nn.Parameter to have static shapes even with symbolic values (#2749)

Author: kshitij12345

thunder/core/proxies.py

@@ -2009,7 +2009,8 @@ def tensorproxy(t: torch.Tensor, /, *, name: None | str, history: None | tuple =
     # See Note [DistributedDataParallel and distparallel_type]
     distparallel_type = getattr(t, "distparallel_type", None)
     _thunder_fsdp_padding_size = getattr(t, "_thunder_fsdp_padding_size", None)
-    if using_symbolic_values():
+    # 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])
         shape = tuple(
             IntegerProxy(

thunder/tests/test_jit_general.py

@@ -1581,6 +1581,24 @@ def foo(a, v):
     assert_close(actual, expected)
 
 
+def test_cache_symbolic_values_nn_parameter_static_shape():
+    linear = torch.nn.Linear(2, 2)
+    x = torch.randn(2, 2)
+
+    jlinear = thunder_jit(linear, cache="symbolic values")
+
+    jlinear(x)
+    exec_trc = thunder.last_traces(jlinear)[-1]
+    for bsym in exec_trc.bound_symbols:
+        if bsym.sym.name == prims.PrimIDs.UNPACK_TRIVIAL and "weight" in bsym.output.name:
+            assert bsym.output.shape == (2, 2)
+        elif bsym.sym.name == prims.PrimIDs.UNPACK_TRIVIAL and "bias" in bsym.output.name:
+            assert bsym.output.shape == (2,)
+        elif bsym.sym.name == prims.PrimIDs.UNPACK_TRIVIAL:  # Input TensorProxy, this should have symbolic values
+            assert isinstance(bsym.output.shape[0], thunder.core.proxies.IntegerProxy)
+            assert isinstance(bsym.output.shape[1], thunder.core.proxies.IntegerProxy)
+
+
 def test_specific_dataclass_returns():
     import transformers