Lightning-AI / Lightning-AI/lightning-thunder

Review what should be returned from ops that return view tensors

Open
#2,216 0 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

There are some Symbols that return a new tensor while the corresponding PyTorch ops return a view.
The set of such ops is defined here -- https://github.com/Lightning-AI/lightning-thunder/blob/62d9535bfeab621b571c3cc8aed974b64381aad0/thunder/torch/__init__.py#L6634-L6660.

For those ops, it might be beneficial to update TensorProxy so that we can tell the outputs of such ops share underlying storage with some other TensorProxy's, then it'd be reasonable to expect improved support for in-place ops on those view tensors.

Status Quo: Returning a new tensor from those ops isn't 100% correct but in-place functionalization and alias updates for in-place both special-case those ops -- https://github.com/search?q=repo%3ALightning-AI%2Flightning-thunder+syms_returning_views&type=code


This issue derives from Mike's review on #2213:

Won't this be incorrect with inplace operations, because this will create a copy, but if the result of the view is modified inplace then the original tensor should be changed?

Originally posted by @mruberry in https://github.com/Lightning-AI/lightning-thunder/pull/2213#discussion_r2138237788

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

Start with the syms_returning_views definition in thunder/torch/init.py around lines 6634-6660, then inspect the in-place functionalization and alias-update code found by the repository search. Trace how TensorProxy represents outputs from these operations and determine what changes are needed for view outputs to share storage correctly. Done means the affected operations have correct view semantics and improved in-place support, with the relevant behavior covered by tests if existing tests are identified.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.