Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

jaxlib 0.4.9 broken on Mac ARM #15951

Closed
hawkinsp opened this issue May 10, 2023 · 4 comments
Closed

jaxlib 0.4.9 broken on Mac ARM #15951

hawkinsp opened this issue May 10, 2023 · 4 comments
Labels
bug Something isn't working

Comments

@hawkinsp
Copy link
Collaborator

hawkinsp commented May 10, 2023 •
edited
Loading

The released jaxlib 0.4.9 wheel for Mac ARM appears to be broken.

In [1]: import jax

In [2]: jax.devices()
Out[2]: [CpuDevice(id=0)]

In [3]: jax.lax.add(1, 2)
'apple-m1' is not a recognized processor for this target (ignoring processor)
'apple-m1' is not a recognized processor for this target (ignoring processor)
'apple-m1' is not a recognized processor for this target (ignoring processor)
'apple-m1' is not a recognized processor for this target (ignoring processor)
'apple-m1' is not a recognized processor for this target (ignoring processor)
'apple-m1' is not a recognized processor for this target (ignoring processor)
'apple-m1' is not a recognized processor for this target (ignoring processor)
'apple-m1' is not a recognized processor for this target (ignoring processor)
LLVM ERROR: 64-bit code requested on a subtarget that doesn't support it!
Abort trap: 6

We have no CI testing for Mac ARM wheels, because none of our CI clusters have such machines. (We build the wheels by cross-compilation, but we cannot test them.)

The workaround for the moment is to downgrade jaxlib to 0.4.7, which was the previous jaxlib release. jaxlib 0.4.7 is compatible with jax 0.4.9.
@hawkinsp hawkinsp added the bug Something isn't working label May 10, 2023
@hawkinsp
Copy link
Collaborator Author

The wheel builds fine if built from source on a Mac ARM host machine, so the issue is probably related to misdetection of the target architecture during cross-compilation.

openxla/xla#2013 may have been the cause.

copybara-service bot pushed a commit to google/tsl that referenced this issue May 11, 2023
@hawkinsp @copybara-github
[LLVM] Restore LLVM toolchains patch removed by openxla/xla#2013
391748a 
This patch turns out to be essential for Mac x86->Mac ARM cross compilation to work (jax-ml/jax#15951).

Unfortunately there's no way to test this except manually, because we have no Mac ARM machines in CI at the moment.

PiperOrigin-RevId: 531065713
copybara-service bot pushed a commit to openxla/xla that referenced this issue May 11, 2023
@hawkinsp @copybara-github
[LLVM] Restore LLVM toolchains patch removed by #2013
1d4e70b 
This patch turns out to be essential for Mac x86->Mac ARM cross compilation to work (jax-ml/jax#15951).

Unfortunately there's no way to test this except manually, because we have no Mac ARM machines in CI at the moment.

PiperOrigin-RevId: 531065713
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this issue May 11, 2023
@hawkinsp @tensorflower-gardener
[LLVM] Restore LLVM toolchains patch removed by openxla/xla#2013
a67243e 
This patch turns out to be essential for Mac x86->Mac ARM cross compilation to work (jax-ml/jax#15951).

Unfortunately there's no way to test this except manually, because we have no Mac ARM machines in CI at the moment.

PiperOrigin-RevId: 531065713
@barlowtwin
Copy link

barlowtwin commented May 11, 2023 •
edited
Loading

Downgrading to jaxlib == 0.4.7 works on M2

@hawkinsp
Copy link
Collaborator Author

openxla/xla@1d4e70b should have fixed the problem, we'll try to release 0.4.10 with a fix.

@hawkinsp
Copy link
Collaborator Author

This was fixed by 0.4.10! Sorry for the breakage. Hopefully we'll get ARM Macs in our CI soon.

@anstadnik anstadnik mentioned this issue May 17, 2023
Merged
3 tasks
@eugene eugene mentioned this issue May 24, 2023
Closed
8 tasks
@thomaspinder thomaspinder mentioned this issue Aug 4, 2023
Closed
Thomas-Christie added a commit to Thomas-Christie/GPJax that referenced this issue Aug 25, 2023
@Thomas-Christie
Bump jax and jaxlib versions
4e18cc4 
In response to issue (jax-ml/jax#15951)
causing us to pin `jaxlib` being fixed.
@Thomas-Christie Thomas-Christie mentioned this issue Aug 25, 2023
Merged
5 tasks
thomaspinder pushed a commit to JaxGaussianProcesses/GPJax that referenced this issue Aug 25, 2023
@Thomas-Christie
Bump jax and jaxlib versions (#366)
c5eb47b 
* Bump `jax` and `jaxlib` versions

In response to issue (jax-ml/jax#15951)
causing us to pin `jaxlib` being fixed.

* Bump `orbax-checkpoint` version
jayfurmanek pushed a commit to ROCm/tensorflow-upstream that referenced this issue Oct 8, 2024
@hawkinsp @jayfurmanek
[LLVM] Restore LLVM toolchains patch removed by openxla/xla#2013
fd11e78 
This patch turns out to be essential for Mac x86->Mac ARM cross compilation to work (jax-ml/jax#15951).

Unfortunately there's no way to test this except manually, because we have no Mac ARM machines in CI at the moment.

PiperOrigin-RevId: 531065713
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
bug Something isn't working
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
2 participants
@hawkinsp @barlowtwin