Lightning-AI / Lightning-AI/lightning-thunder
[inplace] Attempting to reshape a.shape=(3, 3) to shape=(3,)
Open
Nobody has claimed this yet.
bug
in-place
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
Found while creating the smaller repro for #1235
import thunder
import torch
def foo(x, y):
return (x + y).add_(0)
x = torch.randn(3, 3)
foo(x, x[0])
thunder.jit(foo)(x, x[0])
Error
Traceback (most recent call last):
File "lightning-thunder/scratchpad/test_inplace.py", line 44, in <module>
thunder.jit(foo)(x, x[0])
File "lightning-thunder/thunder/__init__.py", line 717, in fn_
cache_entry, inps, pro_to_epi = get_computation_and_inputs(*args, **kwargs)
File "lightning-thunder/thunder/core/langctxs.py", line 136, in _fn
result = fn(*args, **kwargs)
File "lightning-thunder/thunder/__init__.py", line 219, in cache_info_wrapper
res = fn(*args, **kwargs)
File "lightning-thunder/thunder/__init__.py", line 530, in get_computation_and_inputs
functionalize_inplace_ops(
File "lightning-thunder/thunder/core/functionalization.py", line 885, in functionalize_inplace_ops
no_implicit_alias_trace, swap_map_for_aliases = replace_args_with_alias_map(computation_trace, alias_tensor_indices)
File "lightning-thunder/thunder/core/functionalization.py", line 172, in replace_args_with_alias_map
reshaped_arg = prims.reshape.meta(arg, arg_to_replace.shape)
File "lightning-thunder/thunder/core/langctxs.py", line 136, in _fn
result = fn(*args, **kwargs)
File "lightning-thunder/thunder/core/prims.py", line 3167, in reshape_meta
utils.check(
File "lightning-thunder/thunder/core/baseutils.py", line 107, in check
raise exception_type(s())
RuntimeError: Attempting to reshape a.shape=(3, 3) to shape=(3,), but a.numel=9 is different from the number of elements in shape, 3
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Reproduce the failure from scratchpad/test_inplace.py, then inspect thunder/core/functionalization.py at replace_args_with_alias_map and thunder/core/prims.py at reshape_meta. Done means thunder.jit(foo)(x, x[0]) completes without the reshape error and preserves the eager call's behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100