Fix KeyError 'i23' with symbolic shapes and reshape bsyms (#2764)

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

Author: 3 people

thunder/core/update_aliases.py

@@ -90,21 +90,23 @@ def replace_args_with_alias_map(
             reshaped_arg = arg
             if arg_to_replace.shape != arg.shape:
                 with tracectx(computation_trace):
-                    reshaped_arg = prims.reshape.meta(arg, arg_to_replace.shape)
-                    arg_to_optional_bsyms[variableify(arg_to_replace)] = prims.reshape.bind(
-                        arg,
-                        arg_to_replace.shape,
-                        output=reshaped_arg,
-                    )
+                    shape = prims.shape.meta(arg_to_replace)
+                    reshaped_arg = prims.reshape.meta(arg, shape)
+                    reshape_bsym = prims.reshape.bind(arg, shape, output=reshaped_arg)
+                    if using_symbolic_values():
+                        shape_bsym = prims.shape.bind(arg_to_replace, output=shape)
+                        arg_to_optional_bsyms[variableify(arg_to_replace)] = (shape_bsym, reshape_bsym)
+                    else:
+                        arg_to_optional_bsyms[variableify(arg_to_replace)] = (reshape_bsym,)
             swap_map_for_aliases[variableify(arg_to_replace)] = reshaped_arg
     appended_bsyms = {}
     for bsym in computation_trace.bound_symbols:
         for arg in filter(lambda p: isinstance(p, TensorProxy), bsym.flat_args):
-            reshape_bsym = arg_to_optional_bsyms.get(variableify(arg))
-            if reshape_bsym is not None:
-                if reshape_bsym not in appended_bsyms:
-                    bsyms.append(reshape_bsym)
-                    appended_bsyms[reshape_bsym] = arg
+            reshape_bsyms = arg_to_optional_bsyms.get(variableify(arg))
+            if reshape_bsyms is not None:
+                if reshape_bsyms not in appended_bsyms:
+                    bsyms.extend(reshape_bsyms)
+                    appended_bsyms[reshape_bsyms] = arg
         if replaced_args_map := {
             x.name: swap_map_for_aliases[variableify(x)].name
             for x in filter(lambda p: isinstance(p, TensorProxy), bsym.flat_args)

thunder/tests/test_update_aliases.py

@@ -328,8 +328,9 @@ def g(a, b):
 
 @instantiate(
     dtypes=NOTHING,
+    decorators=(pytest.mark.parametrize("cache", ("constant values", "symbolic values")),),
 )
-def test_aliased_input(executor, device, dtype):
+def test_aliased_input(executor, device, dtype, cache):
     def f(x, y, z):
         return y.exp_().add(x) + z.exp()
 
@@ -339,7 +340,7 @@ def f(x, y, z):
     a_ = a.clone().detach()
     b_ = b.clone().detach()
     c_ = c.clone().detach()
-    jfn = executor.make_callable(f)
+    jfn = executor.make_callable(f, cache=cache)
     actual = jfn(a, b, c)
     expected = f(a_, b_, c_)
     torch.testing.assert_close(actual, expected)
@@ -350,8 +351,9 @@ def f(x, y, z):
 
 @instantiate(
     dtypes=NOTHING,
+    decorators=(pytest.mark.parametrize("cache", ("constant values", "symbolic values")),),
 )
-def test_write_to_intermediate_result(executor, device, dtype):
+def test_write_to_intermediate_result(executor, device, dtype, cache):
     if executor == nvFuserExecutor:
         pytest.xfail("nvFuser does not support writing to intermediate results")
 
@@ -361,7 +363,7 @@ def fn(x):
         return y
 
     a = make_tensor((2, 3), dtype=torch.float32, device=device)
-    jfn = executor.make_callable(fn, skip_inplace_alias_updates=True)
+    jfn = executor.make_callable(fn, cache=cache)
     actual = jfn(a)
     expected = fn(a)
     torch.testing.assert_close(actual, expected)
@@ -521,8 +523,9 @@ def foo(x):
 
 @instantiate(
     dtypes=(dtypes.float32,),
+    decorators=(pytest.mark.parametrize("cache", ("constant values", "symbolic values")),),
 )
-def test_aliasing_for_viewed_input_of_different_shapes(executor, device, dtype):
+def test_aliasing_for_viewed_input_of_different_shapes(executor, device, dtype, cache):
     def f(x, y, z):
         return x + 2, y.add_(z)
 
@@ -532,7 +535,7 @@ def f(x, y, z):
     a_ = a.clone().detach()
     b_ = a_[0, :]
     c_ = a_[1, :]
-    jfn = executor.make_callable(f)
+    jfn = executor.make_callable(f, cache=cache)
     actual = jfn(a, b, c)
     expected = f(a_, b_, c_)
     torch.testing.assert_close(actual, expected)