Multi-reviewer summary (readability, code, critical, deep-semantic, integration)
Both lowerings were verified correct on their covered paths. bincount's scatter-add (O(N + depth) vs dense one-hot) and its empty-input Concat([0]) reduce-guard are nicely done; the bool-mask index_put rewrite's NonZero row-major ordering and ScatterND index/update shape rules check out, and the signature change (Sequence[Optional[Union[INT64, BOOL]]]) closes a real interface-drift gap with the caller. accumulate=add matches the existing advanced-index idiom.
Major
- Multi-mask
index_put assumes equal True counts (flagged by 3 reviewers). Concat-ing per-mask NonZero positions requires identical lengths, but torch broadcasts a count-1 mask against a count-k mask, so e.g. mask0 with 1 True and mask1 with 3 True is valid in torch but fails here. This is already noted as a known limitation in the PR description. A fix would broadcast the per-mask position arrays to the common advanced-index shape (or route through the integer-index path) before building the ScatterND index. Note the multi-mask test only covers the equal-count case.
Minor
bincount runtime contracts. The exported model accepts runtime inputs: negatives become negative ScatterElements indices (-1 silently increments the last bucket), and depth = max + 1 can overflow at INT64_MAX. The comment acknowledges torch rejects negatives, but there's no runtime guard. Edge cases / trust-boundary note.
- Bool dispatch is broad.
aten_index_put routes to _aten_index_put_bool if any index is bool, but the helper rejects None/mixed indices, so x[:, mask] = v hard-fails rather than being handled or cleanly decomposed.
Readability
advanced_indices (a list of unsqueezed coordinate tensors) collides with the identically-named list[int] of dimension positions in the adjacent aten_index_put; rename to coord_tensors.
minus_one vs neg_1 are the same constant under two names within one function — neg_1 is the file-wide convention.
The index → bool_mask rename and the same_shape Optional[ir.Shape] None-guard are good clarity/correctness wins.
Originally posted by @titaiwangms in #2921 (comment)
Multi-reviewer summary (readability, code, critical, deep-semantic, integration)
Both lowerings were verified correct on their covered paths.
bincount's scatter-add (O(N + depth) vs dense one-hot) and its empty-inputConcat([0])reduce-guard are nicely done; the bool-maskindex_putrewrite'sNonZerorow-major ordering andScatterNDindex/update shape rules check out, and the signature change (Sequence[Optional[Union[INT64, BOOL]]]) closes a real interface-drift gap with the caller.accumulate=addmatches the existing advanced-index idiom.Major
index_putassumes equal True counts (flagged by 3 reviewers).Concat-ing per-maskNonZeropositions requires identical lengths, but torch broadcasts a count-1 mask against a count-k mask, so e.g.mask0with 1 True andmask1with 3 True is valid in torch but fails here. This is already noted as a known limitation in the PR description. A fix would broadcast the per-mask position arrays to the common advanced-index shape (or route through the integer-index path) before building theScatterNDindex. Note the multi-mask test only covers the equal-count case.Minor
bincountruntime contracts. The exported model accepts runtime inputs: negatives become negativeScatterElementsindices (-1silently increments the last bucket), anddepth = max + 1can overflow atINT64_MAX. The comment acknowledges torch rejects negatives, but there's no runtime guard. Edge cases / trust-boundary note.aten_index_putroutes to_aten_index_put_boolif any index is bool, but the helper rejectsNone/mixed indices, sox[:, mask] = vhard-fails rather than being handled or cleanly decomposed.Readability
advanced_indices(a list of unsqueezed coordinate tensors) collides with the identically-namedlist[int]of dimension positions in the adjacentaten_index_put; rename tocoord_tensors.minus_onevsneg_1are the same constant under two names within one function —neg_1is the file-wide convention.The
index→bool_maskrename and thesame_shapeOptional[ir.Shape]None-guard are good clarity/correctness wins.Originally posted by @titaiwangms in #2921 (comment)