torch.export + torch.onnx.export(dynamo=True) gives incorrect results for scatter_reduce_(reduce="mean")
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.6k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
**Observed behaviour:**
- `scatter_reduce(mean)` shows a large mismatch after export:
- `max_abs_diff: 10.0`
- `mean_abs_diff: 5.5`
- Equivalent `sum/count` control matches exactly:
- `max_abs_diff: 0.0`
- `mean_abs_diff: 0.0`
**Expected behaviour:**
- ONNX output should match eager PyTorch semantics (within normal floating-point tolerance) for:
- `scatter_reduce_(reduce="mean", include_self=False)`
**Impact:**
- Silent numerical correctness issue (wrong predictions without a crash).
**Environment:**
- `torch`: `2.7.1`
- `onnxruntime`: `1.24.1`
- `python`: `3.12`
- `os`: `macOS (darwin 24.6.0)`
**Code example:**
```
import numpy as np
import onnxruntime as ort
import torch
class ScatterMeanModel(torch.nn.Module):
def forward(self, h: torch.Tensor, batch: torch.Tensor) -> torch.Tensor:
# h: [N, F], batch: [N] with group IDs in [0, G-1]
index = batch.unsqueeze(1).repeat(1, h.shape[1])
groups = batch.max().int() + 1
out = torch.zeros(groups, h.shape[1], dtype=h.dtype, device=h.device)
out = out.scatter_reduce_(0, index, h, reduce="mean", include_self=False)
return out
class ScatterSumDivCountModel(torch.nn.Module):
# Mathematically equivalent grouped mean = sum / count.
def forward(self, h: torch.Tensor, batch: torch.Tensor) -> torch.Tensor:
index = batch.unsqueeze(1).repeat(1, h.shape[1])
groups = batch.max().int() + 1
sums = torch.zeros(groups, h.shape[1], dtype=h.dtype, device=h.device)
sums = sums.scatter_reduce_(0, index, h, reduce="sum", include_self=False)
ones = torch.ones(h.shape[0], 1, dtype=h.dtype, device=h.device)
counts = torch.zeros(groups, 1, dtype=h.dtype, device=h.device)
counts = counts.scatter_reduce_(
0, batch.unsqueeze(1), ones, reduce="sum", include_self=False
)
return sums / counts
def run(model: torch.nn.Module) -> tuple[np.ndarray, np.ndarray, float, float]:
model.eval()
h = torch.tensor(
[
[1.0, 10.0],
[3.0, 30.0],
[5.0, 50.0],
[7.0, 70.0],
[2.0, 20.0],
[4.0, 40.0],
],
dtype=torch.float32,
)
batch = torch.tensor([0, 0, 1, 1, 2, 2], dtype=torch.int64)
with torch.inference_mode():
pt = model(h, batch).cpu().numpy()
exported = torch.export.export(model, (h, batch), strict=False)
onnx_program = torch.onnx.export(exported, f=None, dynamo=True)
sess = ort.InferenceSession(
onnx_program.model_proto.SerializeToString(),
providers=["CPUExecutionProvider"],
)
input_names = [i.name for i in sess.get_inputs()]
ort_out = sess.run(
None, {input_names[0]: h.numpy(), input_names[1]: batch.numpy()}
)[0]
diff = np.abs(pt - ort_out)
return pt, ort_out, float(diff.max()), float(diff.mean())
print("torch:", torch.__version__)
print("onnxruntime:", ort.__version__)
pt, ort_out, max_abs, mean_abs = run(ScatterMeanModel())
print("\n=== scatter_reduce(mean) ===")
print("PyTorch output:\n", pt)
print("ONNX Runtime output:\n", ort_out)
print("max_abs_diff:", max_abs)
print("mean_abs_diff:", mean_abs)
pt2, ort_out2, max_abs2, mean_abs2 = run(ScatterSumDivCountModel())
print("\n=== sum/count control ===")
print("PyTorch output:\n", pt2)
print("ONNX Runtime output:\n", ort_out2)
print("max_abs_diff:", max_abs2)
print("mean_abs_diff:", mean_abs2)
```
**Example output:**
```
python tmp/repro_pytorch_scatter_reduce_mean_onnx.py
torch: 2.7.1
onnxruntime: 1.24.1
W0225 11:07:22.309000 94867 torch/onnx/_internal/exporter/_registration.py:103] torchvision is not installed. Skipping torchvision::nms
W0225 11:07:22.310000 94867 torch/onnx/_internal/exporter/_registration.py:103] torchvision is not installed. Skipping torchvision::roi_align
W0225 11:07:22.310000 94867 torch/onnx/_internal/exporter/_registration.py:103] torchvision is not installed. Skipping torchvision::roi_pool
[torch.onnx] Run decomposition...
[torch.onnx] Run decomposition... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
=== scatter_reduce(mean) ===
PyTorch output:
[[ 2. 20.]
[ 6. 60.]
[ 3. 30.]]
ONNX Runtime output:
[[ 3. 30.]
[ 7. 70.]
[ 4. 40.]]
max_abs_diff: 10.0
mean_abs_diff: 5.5
W0225 11:07:22.633000 94867 torch/onnx/_internal/exporter/_registration.py:103] torchvision is not installed. Skipping torchvision::nms
W0225 11:07:22.634000 94867 torch/onnx/_internal/exporter/_registration.py:103] torchvision is not installed. Skipping torchvision::roi_align
W0225 11:07:22.634000 94867 torch/onnx/_internal/exporter/_registration.py:103] torchvision is not installed. Skipping torchvision::roi_pool
[torch.onnx] Run decomposition...
[torch.onnx] Run decomposition... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
=== sum/count control ===
PyTorch output:
[[ 2. 20.]
[ 6. 60.]
[ 3. 30.]]
ONNX Runtime output:
[[ 2. 20.]
[ 6. 60.]
[ 3. 30.]]
max_abs_diff: 0.0
mean_abs_diff: 0.0
```
### Versions
Collecting environment information...
PyTorch version: 2.7.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A
OS: macOS 15.7.3 (arm64)
GCC version: Could not collect
Clang version: 17.0.0 (clang-1700.0.13.5)
CMake version: Could not collect
Libc version: N/A
Python version: 3.12.9 (main, Mar 17 2025, 21:36:21) [Clang 20.1.0 ] (64-bit runtime)
Python platform: macOS-15.7.3-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A
CPU:
Apple M4
Versions of relevant libraries:
[pip3] Could not collect
[conda] Could not collect
______________________
onnxruntime version used in repro: 1.24.1
cc @justinchuby @titaiwangms
Contributor guide
Assessment
This issue has not been assessed yet.