Fix argmin flattened reduction on Ascend (tie break, block size, keepdim) - #5430
Open
zheng1 wants to merge 1 commit into
Open
Fix argmin flattened reduction on Ascend (tie break, block size, keepdim)#5430zheng1 wants to merge 1 commit into
zheng1 wants to merge 1 commit into
Conversation
zheng1
requested review from
0x45f,
103yiran,
Caeruleann,
bin913,
douxetpur,
hellojack163,
huangyiqun,
tengqm and
w1120029931-bit
as code owners
August 13, 2026 00:10
argmin(inp) with dim=None returned wrong indices on the Ascend backend. Three issues on the same code path: 1. Tie breaking. tl.min(..., return_indices=True) did not request return_indices_tie_break_left, so equal values could resolve to an arbitrary index instead of the first one. PyTorch returns the first occurrence. 2. Block size. block_size was derived only from sqrt(numel) with no upper bound. For large inputs it exceeded what the device could handle on this path and produced wrong indices. It is now capped at 1024. 3. keepdim. The output was allocated with torch.empty([]) unconditionally, so keepdim=True was dropped. PyTorch returns shape [1] * inp.dim(). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
zheng1
force-pushed
the
fix/ascend-argmin-reduction
branch
from
August 13, 2026 04:52
8e5e384 to
5508c80
Compare
zheng1
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR Category
Operator
Type of Change
Bug Fix
Description
argmin(inp)withdim=Nonereturned wrong indices on the Ascend backend.Three issues on the same code path:
tl.min(..., return_indices=True)did not requestreturn_indices_tie_break_left, so equal values could resolve to an arbitraryindex instead of the first one. PyTorch returns the first occurrence.
block_sizewas derived only fromsqrt(numel)with no upperbound. For large inputs it exceeded what the device could handle on this path
and produced wrong indices. It is now capped at 1024.
torch.empty([])unconditionally, sokeepdim=Truewas dropped. PyTorch returns shape[1] * inp.dim().Two regression tests cover the large flattened index and NaN handling.
The non-contiguous input problem on the same code path is handled separately in
#5439, which covers it across all the affected reduction operators rather than
just
argmin.Issue
None.
Progress
Performance
The block size cap keeps the reduction within a supported launch configuration.
Tie breaking and the output shape have no performance impact.