Skip to content

Fix fast.cross_entropy VJP on non-contiguous Metal inputs - #4702

Open
StockerMC wants to merge 1 commit into
ml-explore:mainfrom
StockerMC:fix-cross-entropy-vjp-strides
Open

StockerMC wants to merge 1 commit into
ml-explore:mainfrom
StockerMC:fix-cross-entropy-vjp-strides

Conversation

@StockerMC

@StockerMC StockerMC commented Oct 15, 2026 •

Copy link
Copy Markdown
  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: I used Claude Code to understand the relevant code paths, help verify the root cause, and write the fix and test. I reviewed, rebuilt and tested everything myself.

fixes #4686

On Metal, the gradient of mx.fast.cross_entropy can be wrong for a non-row-contiguous view like logits[:, :-1] once the input is big enough that the command buffer gets committed right after the VJP. It's then possible for the next buffer to read the gradient before it's written, so calls sometimes fail.

As was mentioned in the issue, ensure_row_contiguous registers the input copy as a temporary, and the VJP then donates that copy to the output. end_encoding drops temporaries from the encoder's outputs, so the gradient buffer isn't tracked by the fence and the next command buffer doesn't wait for it to be written.

This fix no longer registers the input copy as a temporary, since the output shares its buffer, matching the CUDA VJP and LayerNormVJP.

I added a regression test with an input big enough so that the command buffer always splits after the VJP, and I built and tested on an M3 Pro.

@utsav1033 utsav1033 left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

the script showed that it computed gradient thrice and printed its size, correct being 34.352, which i got right everytime. and that wasnt the case on main as it was 0.0 at first and 6218.577 later.
the new test proved through a test added by this pr, the bug test failed on main and passed on this pr.
nothing else broke when i ran the related test files(test_fast, test_losses, test_autograd, test_vmap)
my setup is apple m4 on macOS 27.0.1

@StockerMC
StockerMC force-pushed the fix-cross-entropy-vjp-strides branch from 8514faa to 46087d0 Compare October 19, 2026 17:35

This branch has not been deployed

No deployments
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.

[BUG] fast.cross_entropy VJP returns wrong gradients for large non-contiguous inputs on Metal (0.32.3 regression)

2 participants