Handle broadcasted bool masks in aten.index_put

[Copilot]

aten.index_put's bool-mask lowering handled the covered single-mask path correctly, but the multi-mask path assumed equal True counts and failed on valid broadcasted cases. It also rejected mixed bool/None indexing such as x[:, mask] = value, even though the generic advanced-index path can represent that form.

• What changed

  • Restrict the dedicated bool-mask lowering to the single [mask] case.
  • Convert bool indices in mixed / multi-index cases into integer position tensors via NonZero, then reuse the existing advanced-index ScatterND lowering.
  • Make advanced-index rank bookkeeping tolerant of these converted bool indices, whose symbolic shape metadata may be absent during export.

• Behavioral impact

  • Multi-mask bool indexing now supports PyTorch-valid broadcasted cases where one mask selects a single position and another selects multiple positions.
  • Mixed indexing with None and a 1-D bool mask now lowers through the generic path instead of failing early.

• Regression coverage

  • Add an e2e case for broadcasted multi-mask bool indexing.
  • Add an e2e case for bool mask indexing mixed with None.

Example of a newly covered case:

x = torch.zeros((3, 4), dtype=torch.float32)
mask0 = torch.tensor([False, True, False], dtype=torch.bool)
mask1 = torch.tensor([True, False, True, True], dtype=torch.bool)
update = torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32)

torch.ops.aten.index_put(x, [mask0, mask1], update)
# updates x[1, [0, 2, 3]]

[Copilot]

Pull request overview

This PR fixes aten.index_put lowering for boolean masks by avoiding incorrect assumptions in the multi-index bool-mask path and routing mixed/multi boolean indexing through the existing advanced-index (ScatterND) lowering.

Changes:

• Restricts the dedicated boolean-mask lowering to the single-index form ([mask]).
• Converts 1-D boolean indices appearing in mixed/multi-index cases into integer position tensors via NonZero, then reuses the advanced-index ScatterND lowering (including broadcast handling).
• Adds end-to-end export coverage for broadcasted multi-mask boolean indexing and for boolean indexing mixed with None (slice-like) indices.

Reviewed changes

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

File	Description
tests/function_libs/torch_lib/e2e_ops_tests.py	Adds new e2e export test cases covering broadcasted multi-mask bool indexing and bool mask mixed with None.
onnxscript/function_libs/torch_lib/ops/core.py	Updates aten_index_put lowering to route mixed/multi bool masks through the advanced-index path by converting bool masks to integer indices, and adjusts rank bookkeeping accordingly.

[codecov]

Codecov Report

❌ Patch coverage is 62.96296% with 20 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.66%. Comparing base (00de62e) to head (6c844d3).
⚠️ Report is 4 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	62.96%	16 Missing and 4 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2966      +/-   ##
==========================================
+ Coverage   72.64%   72.66%   +0.01%     
==========================================
  Files         265      265              
  Lines       32192    32219      +27     
  Branches     3038     3044       +6     
==========================================
+ Hits        23385    23411      +26     
+ Misses       7776     7774       -2     
- Partials     1031     1034       +3     

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

[Copilot]

Pull request overview

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

[Justin Chu (justinchuby)]

I think this should be fixed, instead