torch.cumprod returns an incorrect nested JVP at a zero factor
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
`torch.cumprod` returns `68.0` for a nested forward-mode second derivative whose mathematically correct value is `144.0` in `float64`.
The cumulative products and the weighted scalar reduction are polynomials. A zero factor is not a nondifferentiable point. At `t = 1.0` the input contains a zero, and expanding the prefix products gives second derivative `144.0`. PyTorch’s nested JVP is `68.0`. 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((t, t + 1, 2 - t, t - 1, 2 * t + 1, t + 3)).reshape(2, 3)
w = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=t.dtype)
return (torch.cumprod(x, dim=-1, dtype=t.dtype) * w).sum()
def jvp1(t):
return torch.func.jvp(f, (t,), (torch.ones_like(t),))[1]
t = torch.tensor(1.0, dtype=dtype)
_, actual = torch.func.jvp(jvp1, (t,), (torch.ones_like(t),))
expected = torch.tensor(144.0, dtype=dtype)
print("actual:")
print(actual)
print("expected:")
print(expected)
```
### Actual vs expected result
```text
actual:
68.0
expected:
144.0
```
The discrepancy occurs in the nested (second-order) forward-mode derivative. The forward value is `11.0` and matches the expanded prefix products.
### 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
```
cc @albanD
Contributor guide
Research direction
Run the provided CPU reproduction and trace the nested torch.func.jvp calls through torch.cumprod. The fix is complete when the second-order forward-mode result is 144.0 for this case and a regression test covers the zero-factor input.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 52/100