FakeTensor returns incorrect output strides for native_batch_norm with noncontiguous CPU input
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
`aten.native_batch_norm` returns a contiguous output for an arbitrary noncontiguous CPU input in eager mode, but FakeTensor preserves the input strides. This causes failure under `torch.compile` and `torch.export`.
Minimal Reproduce
```python
import torch
from torch._subclasses.fake_tensor import FakeTensorMode
x = torch.randn(2, 3, 4).transpose(1, 2)
weight, bias = torch.ones(4), torch.zeros(4)
mean, var = torch.zeros(4), torch.ones(4)
args = (x, weight, bias, mean, var, True, 0.1, 1e-5)
eager = torch.ops.aten.native_batch_norm.default(*args)[0]
with FakeTensorMode() as mode:
fake_args = tuple(mode.from_tensor(t) if isinstance(t, torch.Tensor) else t for t in args)
fake = torch.ops.aten.native_batch_norm.default(*fake_args)[0]
print("input:", x.shape, x.stride())
print("eager:", eager.shape, eager.stride())
print("fake:", fake.shape, fake.stride())
```
Output
```text
input: torch.Size([2, 4, 3]) (12, 1, 4)
eager: torch.Size([2, 4, 3]) (12, 3, 1)
fake: torch.Size([2, 4, 3]) (12, 1, 4)
```
Expected
```text
input: torch.Size([2, 4, 3]) (12, 1, 4)
eager: torch.Size([2, 4, 3]) (12, 3, 1)
fake: torch.Size([2, 4, 3]) (12, 3, 1)
```
The incorrect metadata breaks
```python
def fn(x, weight, bias, mean, var):
y = torch.ops.aten.native_batch_norm.default(x, weight, bias, mean, var, True, 0.1, 1e-5)[0]
return y.view(-1)
torch.compile(fn, backend="eager", fullgraph=True)(x, weight, bias, mean, var)
```
Error
```text
torch._dynamo.exc.TorchRuntimeError: RuntimeError when making fake tensor call
Explanation: Dynamo failed to run FX node with fake tensors: call_method view(*(FakeTensor(..., size=(2, 4, 3)), -1), **{}): got ValueError('Cannot view a tensor with shape torch.Size([2, 4, 3]) and strides (12, 1, 4) as a tensor with shape (24,)!')
```
Compilation fails during FakeTensor execution because `view(-1)` is evaluated against strides `(12, 1, 4)` instead of the native CPU output strides `(12, 3, 1)`. `torch.export` fails for the same reason. FakeTensor should match the native CPU output layout while preserving supported memory formats such as channels-last.
### Versions
```text
PyTorch version: 2.14.0a0+git4c5b5ab
OS: macOS 14.1 (arm64)
Python version: 3.12.13
CPU: Apple M3 Pro
```
cc @chauhang @penguinwu @eellison @aorenste @liangel-02 @bdhirsh @bobrenjc93
Contributor guide
Research direction
Start by reproducing the minimal example with FakeTensorMode and aten.native_batch_norm.default, then trace the FakeTensor handling used by torch.compile and torch.export. Compare fake and eager output strides for the noncontiguous CPU input while checking that channels-last layouts remain supported. Done means the fake output matches the native CPU layout and both compilation examples succeed.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100