Skip to content

Refactor CPU ScatterAxis and MaskedScatter dtype dispatch - #4176

Merged
zcbenz merged 2 commits into
ml-explore:mainfrom
PhysicistJohn:dispatch-cpu-scatter-axis-main
Aug 11, 2026
Merged

Refactor CPU ScatterAxis and MaskedScatter dtype dispatch#4176
zcbenz merged 2 commits into
ml-explore:mainfrom
PhysicistJohn:dispatch-cpu-scatter-axis-main

Conversation

@PhysicistJohn

@PhysicistJohn PhysicistJohn commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Summary

Replace the remaining exhaustive payload dtype switches in CPU ScatterAxis
and MaskedScatter with the existing dispatch_all_types helper. This keeps
the restricted integer index dispatch, reduction mode selection, broadcasting,
source traversal, and diagnostics unchanged while removing 83 net lines.

Testing

  • Release CPU-only build with global -Werror
  • Native C++ suite (247/247 cases, 3,326/3,326 assertions)
  • python/tests/test_ops.py (153/153 tests)
  • focused vmap and autograd coverage for both operations
  • 1,456 ScatterAxis parity cases and 14 exact unsupported-index diagnostics
  • 602 MaskedScatter parity cases and 14 exact short-source diagnostics
  • baseline/candidate symbol and linked-export sets match; the indexing object is 416 bytes smaller
  • clang-format and git diff --check

@zcbenz

zcbenz commented Aug 11, 2026

Copy link
Copy Markdown
Member

Would you mind updating all the dispatch code in the file in one pull request? Our CI is throttled to handle many pull requests.

@PhysicistJohn PhysicistJohn changed the title Refactor CPU ScatterAxis dtype dispatch Refactor CPU ScatterAxis and MaskedScatter dtype dispatch Aug 11, 2026
@PhysicistJohn

Copy link
Copy Markdown
Contributor Author

Thanks, Cheng — agreed. I folded the remaining MaskedScatter payload switch
into this PR, so the last manual exhaustive payload-dtype router in
indexing.cpp now uses dispatch_all_types. The restricted integral-index
routers and reduction-mode switches remain unchanged because they encode
narrower type and operation semantics, including the existing diagnostics. The
combined diff is +8/-91. The Release -Werror build, 247/247 native cases,
153/153 Python ops tests, and paired ScatterAxis/MaskedScatter parity and
diagnostic checks all pass.

@zcbenz
zcbenz merged commit 9e37e5b into ml-explore:main Aug 11, 2026
28 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants