pytorch / pytorch/pytorch

torch.nn.functional.batch_norm returns an incorrect second-order JVP

Open
#196,699 0 comments 0 reactions 0 assignees View on GitHub
bot-triaged module: correctness (silent) module: forward ad module: norms and normalization release triage triage review
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

`torch.nn.functional.batch_norm` returns `-2.132438339486815` for a nested forward-mode second derivative whose mathematically correct value is `-0.20703015646993195` in `float64`.
The example uses training statistics and `eps = 0.25`. The independent reference computes each channel’s mean, population variance, affine transform, and weighted output directly. The positive epsilon removes zero-variance singularities. At `t = -6.0`, the nested JVP disagrees with that reference. The forward value matches.
This issue reproduces on CPU, so CUDA availability does not affect the result.
```python
import torch
dtype = torch.float64
def f(t):
x = torch.stack([(i + 1) * t + (i % 3 - 1) * t * t + i for i in range(8)]).reshape(2, 2, 2)
weight = torch.tensor([1.5, -0.5], dtype=t.dtype)
bias = torch.tensor([0.25, 1], dtype=t.dtype)
y = torch.nn.functional.batch_norm(
x, None, None, weight=weight, bias=bias, training=True, momentum=0.1, eps=0.25
)
w = torch.arange(1, 9, dtype=t.dtype).reshape(2, 2, 2)
return (y * w).sum()
def jvp1(t):
return torch.func.jvp(f, (t,), (torch.ones_like(t),))[1]
t = torch.tensor(-6.0, dtype=dtype)
_, actual = torch.func.jvp(jvp1, (t,), (torch.ones_like(t),))
expected = torch.tensor(-0.20703015646993195, dtype=dtype)
print("actual:")
print(actual)
print("expected:")
print(expected)
```
### Actual vs expected result
```text
actual:
-2.132438339486815
expected:
-0.20703015646993195
```
The discrepancy occurs in the nested (second-order) forward-mode derivative. The forward value is `36.16877756630648` and matches the independent mean/variance reference.

### Versions

```text
PyTorch version: 2.12.0+cu130
Is debug build: False
CUDA used to build PyTorch: 13.0
ROCM used to build PyTorch: N/A
OS: Ubuntu 24.04.3 LTS (x86_64)
Python version: 3.13.13 | packaged by Anaconda, Inc.
Python platform: Linux x86_64
Is CUDA available: False
GPU models:
GPU 0: NVIDIA RTX 6000 Ada Generation
GPU 1: NVIDIA RTX 6000 Ada Generation
Nvidia driver version: 570.211.01
CPU:
Model name: AMD Ryzen Threadripper PRO 7985WX 64-Cores
CPU(s): 128
Versions of relevant libraries:
[pip3] numpy==2.4.6
[pip3] torch==2.12.0
[pip3] triton==3.7.0
```

Contributor guide

Open the contributing guide

Research direction

Start with the torch.nn.functional.batch_norm entry point and run the provided CPU reproducer using float64, training statistics, and eps=0.25. Trace the nested torch.func.jvp path and compare its second-order result with the independent mean, variance, affine-transform, and weighted-output reference; done means the result matches -0.20703015646993195.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.