Skip to content

Fix wrong results in left_shift/right_shift on overflow - #4694

Merged
zcbenz merged 2 commits into
ml-explore:mainfrom
aaishwarymishra:fix-shift
Oct 29, 2026
Merged

Fix wrong results in left_shift/right_shift on overflow#4694
zcbenz merged 2 commits into
ml-explore:mainfrom
aaishwarymishra:fix-shift

Conversation

@aaishwarymishra

@aaishwarymishra aaishwarymishra commented Oct 14, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor
  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure:
    So currently when the shift is greater than dtype size, for example 32 in case of int 32, the rshift and lshift returns the same value instead of 0 and -1 respectively like how numpy and jax does.

Found in data-apis/array-api-compat#451

I tried to have the upper level fix, but it was not possible so I defaulted to modifying the kernels.

The cpu kernel code was written by AI as I am not that much familiar with cpu simd.

example:

import jax.numpy as jnp
import numpy as np
import mlx.core as mx

a = jnp.array([1, 2, -3], dtype=jnp.int32)


print(a>>32)
print(np.array(a.tolist(), dtype=np.int32)>>32)
print(mx.array(a.tolist(), dtype=mx.int32)>>32)

results:

blanky@Blankys-MacBook-Air flax-exp % uv run test.py
[ 0  0 -1]
[ 0  0 -1]
array([1, 2, -3], dtype=int32)

@aaishwarymishra aaishwarymishra changed the title Fixed lshift and tshift to return 0 for unsigned and -1 for signed value when shift is greater than dtype size Fixed lshift and rshift to return 0 for unsigned and -1 for signed value when shift is greater than dtype size Oct 14, 2026
@aaishwarymishra

aaishwarymishra commented Oct 22, 2026

Copy link
Copy Markdown
Contributor Author

@zcbenz Hi, just checking, its fine if you are busy feel free to review when you get the time :)

@zcbenz

zcbenz commented Oct 22, 2026

Copy link
Copy Markdown
Member

There are some high priority PRs that are quite hard to review, so this will wait for a while.

@zcbenz zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Oct 23, 2026
@aaishwarymishra

aaishwarymishra commented Oct 23, 2026

Copy link
Copy Markdown
Contributor Author

I understand, I wanted to ask will it be ok for me to Open more prs? I don't want to overwhelm you with multiple prs if you have limited bandwidth :)

@zcbenz

zcbenz commented Oct 23, 2026

Copy link
Copy Markdown
Member

If you have some important fixes then yes, otherwise we would prefer waiting for a while for us to consume the backlog.

zcbenz
zcbenz approved these changes Oct 29, 2026
View reviewed changes

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

It seems that the mps backend of pytorch has the same bug. I'm good with this change since it is what every other framework does.

@zcbenz zcbenz removed the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Oct 29, 2026
@zcbenz zcbenz changed the title Fixed lshift and rshift to return 0 for unsigned and -1 for signed value when shift is greater than dtype size Fix wrong results in lshift/rshift on overflow Oct 29, 2026
@zcbenz zcbenz changed the title Fix wrong results in lshift/rshift on overflow Fix wrong results in left_shift/right_shift on overflow Oct 29, 2026
@aaishwarymishra

aaishwarymishra commented Oct 29, 2026

Copy link
Copy Markdown
Contributor Author

It seems that the mps backend of pytorch has the same bug. I'm good with this change since it is what every other framework does.

I can make an issue and try to make pr to make it consistent, I am not sure how it will work though as Pytorch directly uses mps as backend.

@zcbenz
zcbenz merged commit 232eb68 into ml-explore:main Oct 29, 2026
30 checks passed
This was referenced Oct 29, 2026
Closed
Open
@dcvz dcvz mentioned this pull request Oct 31, 2026
Open
This was referenced Nov 3, 2026
Open
Open
Open
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@zcbenz zcbenz zcbenz approved these changes

Assignees

No one assigned

Labels

low priority

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@aaishwarymishra @zcbenz