Skip to content

Fix scan and sort for a zero-size axis - #4340

Open
eyupcanakman wants to merge 1 commit into
ml-explore:mainfrom
eyupcanakman:fix/empty-axis-scan-sort
Open

Fix scan and sort for a zero-size axis#4340
eyupcanakman wants to merge 1 commit into
ml-explore:mainfrom
eyupcanakman:fix/empty-axis-scan-sort

Conversation

@eyupcanakman

Copy link
Copy Markdown
Contributor

Proposed changes

Scan and Sort work out the row count on the host with in.size() / in.shape(axis).
When the scanned or sorted axis has length zero, both operands are zero.
On x86-64 the DIV instruction raises #DE, so mx.cumsum(mx.array([]), axis=0) kills the process with SIGFPE.

Apple silicon never sees it.
AArch64 integer division by zero returns 0 and raises nothing, so the loop bound comes out zero and the empty result is right by accident.

I built main for linux/amd64 in a container and ran cumsum on a shape (0,) array.
It dies with a floating point exception, exit 136.
The same build with the guard exits 0 and prints the empty result.

A USE_UBSAN build on macOS reports the division at mlx/backend/cpu/scan.cpp:170 and mlx/backend/cpu/sort.cpp:124 before the change, and nothing after it.

The fix returns early when the output is empty, in the scan and sort entry points of all three backends.
It is the same out.size() == 0 check SearchSorted already carries.
Eight files is one guard landing in three backends across two ops, plus the two test files.

Partition and ArgPartition divide the same way on CPU and are left alone.
mlx/ops.cpp rejects every kth on a zero-length axis, so those primitives are never built.
On Metal and CUDA they route through the shared sort helper and pick up the guard anyway.

I have no NVIDIA hardware here.
Both .cu files compile clean under nvcc 12.9 for sm_80, but running them needs a real GPU, so that part rides on CI.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

@eyupcanakman
eyupcanakman force-pushed the fix/empty-axis-scan-sort branch from beffa50 to f07f722 Compare August 19, 2026 15:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants