Lightning-AI / Lightning-AI/lightning-thunder

[inplace] Attempting to reshape a.shape=(3, 3) to shape=(3,)

Open
#1,237 1 comment 0 reactions 0 assignees View on GitHub

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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.