Jax for Apple M1 GPU?

[fbarfi]

Is there any hope that Jax will be made compatible with Apple M1 GPU? I have am M1 Max with 32 GPU cores remaining idle, it is a loss not to be able to use them in my data science research. M2 Ultra has double GPU cores (64). Apple users in data sciences feel left behind (:O) . The apple market will in a few years be dominated with such machines. No offense to NVIDIA, but we need diversity. Thanks for any comment or feedback. Of course I do not mean to insult anyone. Just concerned.

[hawkinsp]

See #8074

There's a highly experimental path to do this via IREE -> Vulkan and MoltenVK that works for some simple examples.

[adam-hartshorne]

Looks like it is now available.

https://developer.apple.com/metal/jax/

[femtomc]

Has anyone successfully built jaxlib with the environment / build script mentioned in this post?

[fbarfi]

Yes I did on 2023 MacBook Pro M2 machine.

[femtomc]

Can you share what version of clang (and lld) you used to successfully build?

[femtomc]

I'm attempting to do the same, on the exact same device specs:

clang: error: invalid linker name in argument '-fuse-ld=-debugger-tuning=lldb'
Target //build:build_wheel failed to build
INFO: Elapsed time: 124.272s, Critical Path: 13.53s
INFO: 544 processes: 14 internal, 530 local.
FAILED: Build did NOT complete successfully
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully
b''

What is this error indicative of in the build process?

Environment:

jax (21fc6e0) via 🐍 v3.9.12 on ☁️  (us-east-2) took 2m5s
❯ clang --version
Apple clang version 14.0.3 (clang-1403.0.22.14.1)
Target: arm64-apple-darwin22.5.0
Thread model: posix
InstalledDir: /Library/Developer/CommandLineTools/usr/bin

jax (21fc6e0) via 🐍 v3.9.12 on ☁️  (us-east-2)
❯ lldb --version
lldb-1403.0.17.67
Apple Swift version 5.8.1 (swiftlang-5.8.0.124.5 clang-1403.0.22.11.100)

[femtomc]

@hawkinsp if you happen to see this, is there anything stupid I could be doing here? (I'm not sure how much you were involved in jax-metal -- I really don't know what avenues I have to ask the developers questions)

[fbarfi]

Apple clang version 14.0.3 (clang-1403.0.22.14.1)
Target: arm64-apple-darwin22.5.0
Thread model: posix

[fbarfi]

I followed exactly what the following link says step by step:

https://developer.apple.com/metal/jax/

including building a specific jaxlib compatible with metal plug-in