Fix aten_masked_scatter for broadcasting masks

[gonultasbu]

aten_masked_scatter returns wrong results when mask is broadcastable to self but has the same rank (e.g. mask (1, S, 1) vs self (1, S, D)): the rank-only check expands self and leaves mask un-expanded, so NonZero(mask) enumerates only a subset of masked positions and ScatterND writes the wrong source elements (e.g. output [1, 2, 3, 4] instead of [1, 9, 17, 25]).

Fix — expand both to their common shape:

self = op.Expand(self, op.Shape(mask))
mask = op.Expand(mask, op.Shape(self))

The original branch (#2112, the initial add for Gemma3) only handled rank differences and has been unchanged since; the Gemma3 path pre-expands the mask (expand_as), so it was unaffected. Surfaced via torch.onnx.export(dynamo=True) on a model that injects features with masked_scatter — pytorch/pytorch#186146.

Test: adds an ops.aten.masked_scatter OpInfo with broadcastable-mask samples (the existing torch OpInfo didn't cover same-rank broadcasting) — fails on main, passes here.

[gonultasbu]

@microsoft-github-policy-service agree company="Noble Machines"

[Copilot]

Pull request overview

This PR aims to fix incorrect aten_masked_scatter results when mask is broadcastable to self (including same-rank broadcasting), and adds OpInfo-based test coverage to prevent regressions in the torchlib exporter path.

Changes:

• Update aten_masked_scatter to broadcast self/mask before NonZero(mask) so indices cover all masked elements.
• Add a custom ops.aten.masked_scatter OpInfo with broadcastable-mask sample inputs (including same-rank broadcasting).
• Register the new OpInfo in the torchlib ops test matrix.

Reviewed changes

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

File	Description
onnxscript/function_libs/torch_lib/ops/core.py	Adjusts aten_masked_scatter broadcasting logic prior to NonZero/ScatterND.
tests/function_libs/torch_lib/ops_test_data.py	Adds TorchLibOpInfo entry so the new OpInfo is exercised by the torchlib test suite.
tests/function_libs/torch_lib/extra_opinfo.py	Defines ops.aten.masked_scatter OpInfo and sample inputs covering broadcastable masks.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 72.63%. Comparing base (e0d3edc) to head (e1a7d01).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2929      +/-   ##
==========================================
- Coverage   72.63%   72.63%   -0.01%     
==========================================
  Files         259      259              
  Lines       31666    31665       -1     
  Branches     2982     2981       -1     
==========================================
- Hits        23000    22999       -1     
  Misses       7656     7656              
  Partials     1010     1010              

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

[titaiwangms]

Thanks!