Fix scatter_reduce(reduce="mean") producing incorrect ONNX export results

[Copilot]

• Fix aten_scatter_reduce in core.py to implement reduce="mean" correctly
  • include_self=True: sum/count approach including self in both
  • include_self=False: use Where to preserve self[i] for positions with no scattered values (fixes CI failure where original values were incorrectly replaced with 0)
• Fix lint error in e2e_ops_tests.py (forward method signature too long)
• Update xfail in ops_test_data.py to remove general mean xfail
• Add e2e tests for scatter_reduce with reduce="mean" for both include_self values

Original prompt

---

This section details on the original issue you should resolve

<issue_title>torch.export + torch.onnx.export(dynamo=True) gives incorrect results for scatter_reduce_(reduce="mean")</issue_title>
<issue_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
[tor...

</details>



<!-- START COPILOT CODING AGENT SUFFIX -->

- Fixes pytorch/pytorch#175738

<!-- START COPILOT CODING AGENT TIPS -->
---

🔒 GitHub Advanced Security automatically protects Copilot coding agent pull requests. You can protect all pull requests by enabling Advanced Security for your repositories. [Learn more about Advanced Security.](https://gh.io/cca-advanced-security)

[codecov]

Codecov Report

❌ Patch coverage is 55.00000% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.64%. Comparing base (ef2bc22) to head (7c50bd0).
⚠️ Report is 69 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/core.py	55.00%	5 Missing and 4 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2829      +/-   ##
==========================================
+ Coverage   71.78%   72.64%   +0.86%     
==========================================
  Files         239      265      +26     
  Lines       28989    32208    +3219     
  Branches     2859     3041     +182     
==========================================
+ Hits        20809    23398    +2589     
- Misses       7209     7778     +569     
- Partials      971     1032      +61     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

[justinchuby]

@copilot CI errors

[Copilot]

@copilot CI errors

Fixed in 6614c55. There were two CI failures:

1. Test failures (scatter_reduce_mean_cpu_float32): The include_self=False path was returning 0 for positions with no scattered values instead of preserving the original self[i]. Fixed by using Where(count==0, self, mean_vals).

2. Lint failure (RUFF-FORMAT): The forward method signature in e2e_ops_tests.py was too long. Split into multiple lines.

[Copilot]

Pull request overview

This PR fixes ONNX lowering for aten_scatter_reduce(..., reduce="mean") so exported ONNX programs match PyTorch semantics (including correct include_self=False behavior when no values are scattered into a position).

Changes:

• Implement reduce="mean" in aten_scatter_reduce via sum/count (and Where to preserve self values when include_self=False and count==0).
• Update scatter_reduce test expectations (xfail) to remove the blanket "mean" unsupported marker and narrow failures to specific dtype cases.
• Add e2e coverage for scatter_reduce(mean) export for both include_self settings.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
onnxscript/function_libs/torch_lib/ops/core.py	Adds explicit ONNX lowering for scatter_reduce(reduce="mean") using ScatterElements(reduction="add") + division, including include_self=False “preserve self when no updates” handling.
tests/function_libs/torch_lib/ops_test_data.py	Adjusts xfail coverage for scatter_reduce(mean) now that mean lowering exists, scoping remaining failures to specific dtype limitations.
tests/function_libs/torch_lib/e2e_ops_tests.py	Adds new e2e export tests targeting the reported scatter_reduce(mean) issue for both include_self=False and include_self=True.