Skip to content

fix: clamp gather indices in diag 1D→2D for non-zero k#54

Open
ymwangg wants to merge 1 commit intoaws-neuron:mainfrom
ymwangg:fix/diag-oob-indices
Open

fix: clamp gather indices in diag 1D→2D for non-zero k#54
ymwangg wants to merge 1 commit intoaws-neuron:mainfrom
ymwangg:fix/diag-oob-indices

Conversation

@ymwangg
Copy link
Copy Markdown

@ymwangg ymwangg commented Apr 17, 2026

Summary

  • When k != 0, np.diag (1D→2D path) builds an (n, n) output where n = len(v) + |k|, but the iota gather indices range [0, n) while v only has len(v) elements
  • The compiler rejects the out-of-bounds gather even though those positions are masked to zero by where(diag_mask, ...)
  • Fix: clamp idx_ref to [0, len(v) - 1] before take so the gather stays in-bounds

Test plan

  • test_diag_1d_to_2d_negative_k now passes (was failing with NCC_EVRF056 OOB error)
  • All other diag tests (test_diag_1d_to_2d, test_diag_2d_to_1d, test_diagonal_gather) still pass

🤖 Generated with Claude Code

When k != 0, the output matrix is (n, n) where n = len(v) + |k|, but
the iota indices range [0, n) while v only has len(v) elements. The
compiler rejects the out-of-bounds gather even though those positions
are masked to zero. Clamping to [0, len(v)-1] before take fixes this.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@ymwangg ymwangg requested a review from a team April 17, 2026 17:59
Copy link
Copy Markdown
Contributor

@vgene vgene left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

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