Skip to content

Follow up on bincount + index_put #2965

Description

@justinchuby

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 indexbool_mask rename and the same_shape Optional[ir.Shape] None-guard are good clarity/correctness wins.

Originally posted by @titaiwangms in #2921 (comment)

Metadata

Metadata

Labels

No labels
No labels

Type

No type

Fields

No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions