diff --git a/thunder/core/update_aliases.py b/thunder/core/update_aliases.py index 3a679bcf13..de8ec89604 100644 --- a/thunder/core/update_aliases.py +++ b/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) diff --git a/thunder/tests/test_update_aliases.py b/thunder/tests/test_update_aliases.py index e90a463a1a..4d293aae21 100644 --- a/thunder/tests/test_update_aliases.py +++ b/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)