pytorch / pytorch/pytorch

FakeTensor returns incorrect output strides for native_batch_norm with noncontiguous CPU input

Open
#192,668 0 comments 0 reactions 0 assignees View on GitHub
bot-triaged module: fakeTensor module: pt2-dispatcher oncall: pt2 triaged
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.