Skip to content

Fix Metal gather with scalar indices - #3737

Closed
tsnAnh wants to merge 1 commit into
ml-explore:mainfrom
tsnAnh:codex/fix-gather-scalar-index-metal-validation
Closed

tsnAnh wants to merge 1 commit into
ml-explore:mainfrom
tsnAnh:codex/fix-gather-scalar-index-metal-validation

Conversation

@tsnAnh

@tsnAnh tsnAnh commented Jun 21, 2026

Copy link
Copy Markdown

Summary

Fixes a Metal debug validation crash in Gather::eval_gpu when gather/take uses scalar indices.

Gather already has a scalar-index specialization (idx_ndim == 0), but the Metal encoder still receives empty index metadata vectors:

compute_encoder.set_vector_bytes(idx_shapes, 7);
compute_encoder.set_vector_bytes(idx_strides, 8);
compute_encoder.set_vector_bytes(idx_contigs, 9);

With Metal API validation enabled, an empty vector reaches setBytes:length:attributeStride:atIndex: as a nil bytes argument and aborts.

Evidence

Observed downstream crash:

-[MTLDebugComputeCommandEncoder setBytes:length:attributeStride:atIndex:]:387: failed assertion `bytes argument cannot be nil.'
Task: signal SIGABRT

The abort was isolated to this line in mlx/backend/metal/indexing.cpp:

compute_encoder.set_vector_bytes(idx_shapes, 7);

The same file already handles this exact Metal binding edge case in Scatter::eval_gpu by pushing dummy metadata when idx_ndim == 0:

if (idx_ndim == 0) {
  // Add a 0 in idx_shapes and strides to avoid the missing buffer binding
  // error in the metal API.
  idx_shapes.push_back(0);
  idx_strides.push_back(0);
  idx_contigs.push_back(false);
}

This PR applies the same guard to Gather::eval_gpu. idx_ndim remains 0, so the scalar-index specialization behavior does not change; the patch only ensures Metal receives non-nil metadata buffers.

Existing public API coverage already includes scalar-index take:

out = mx.take(a, mx.array(1), axis=0)
self.assertEqual(out.shape, (4,))

Validation

Ran locally:

git diff --check

Downstream validation after applying the same patch in a SwiftPM MLX checkout:

xcodebuild -quiet -project iosApp/iosApp.xcodeproj -scheme iosApp -configuration Debug -destination 'generic/platform=iOS Simulator' -skipMacroValidation -skipPackagePluginValidation build
xcodebuild -quiet -project iosApp/iosApp.xcodeproj -scheme iosApp -configuration Debug -destination 'generic/platform=iOS' -skipMacroValidation -skipPackagePluginValidation build
./gradlew :shared:linkDebugFrameworkIosSimulatorArm64

Those builds passed and no longer hit the Metal nil-buffer assertion.

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.

1 participant