Lightning-AI / Lightning-AI/lightning-thunder
`torch.transpose` seems to be mapped to different ops depending on `requires_grad`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
*Note*: If you have a model or program that is not supported yet but should be, please use the program coverage template.
## 🐛 Bug
I admit I'm not sure if this is a bug of an expected behavior, but `torch.transpose` is mapped to `torch.transpose` if `requires_grad=False`, `torch.permute` otherwise.
### To Reproduce
#### Code sample
```python
import torch
import thunder
@thunder.jit
def f(x):
return x.transpose(0, 1)
for requires_grad in (False, True):
print("#" * 120)
x = torch.rand((4, 2), requires_grad=requires_grad)
f(x)
print(f"### {requires_grad = }")
print(f"### thunder.last_traces(f)[-1]\n{thunder.last_traces(f)[-1]}")
```
#### output
```
########################################################################################################################
### requires_grad = False
### thunder.last_traces(f)[-1]
# Constructed by Unwrap the actual return value
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast
def computation(x):
# x: "cpu f32[4, 2]"
t0 = torch.transpose(x, 0, 1) # t0: "cpu f32[2, 4]"
# t0 = ltorch.transpose(x, 0, 1) # t0: "cpu f32[2, 4]"
# t0 = prims.transpose(x, (1, 0)) # t0: "cpu f32[2, 4]"
return t0
########################################################################################################################
### requires_grad = True
### thunder.last_traces(f)[-1]
# Constructed by Delete Last Used (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast
def computation(x):
# x: "cpu f32[4, 2]"
t0 = torch.permute(x, (1, 0)) # t0: "cpu f32[2, 4]"
# t0 = ltorch.permute(x, (1, 0)) # t0: "cpu f32[2, 4]"
# t0 = prims.transpose(x, (1, 0)) # t0: "cpu f32[2, 4]"
return {'output': t0, 'flat_args': [x], 'flat_output': (t0,)}, ((), ())
```
### Expected behavior
### Environment
- PyTorch Version (e.g., 1.0):
- OS (e.g., Linux):
- How you installed PyTorch (`conda`, `pip`, source):
- Build command you used (if compiling from source):
- Python version:
- CUDA/cuDNN version:
- GPU models and configuration:
- Any other relevant information:
### Additional context
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
Start by running the provided Python reproducer with thunder.jit and inspect thunder.last_traces(f)[-1] for both requires_grad values. Trace the operator-selection path for torch.transpose and determine the intended mapping; done should include a confirmed expected behavior and a regression test covering both cases.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100