JAX acceleration for Mac GPUs.
Project description
jax-metal
The jax-metal package is a Metal GPU plugin to provide Metal acceleration on Mac platforms for JAX applications.
Installation
The following table tracks jax-metal versions and compatible versions of jax, jaxlib and MacOS.
| jax-metal | MacOS | jaxlib | jax |
|---|---|---|---|
| 0.0.6 | Sonoma 14.4 Beta | >=v0.4.22, <v0.4.24 | >=v0.4.22 |
| 0.0.5 | Sonoma 14.2+ | >=v0.4.20, <v0.4.22 | >=v0.4.20 |
| 0.0.4 | Sonoma 14.0+ | v0.4.11 | v0.4.11 |
| 0.0.3 | Ventura 13.4.1+, Sonoma 14.0 Beta | v0.4.10 | v0.4.11 |
We recommend to install the binary package with venv or conda.
python3 -m venv ~/jax-metal
source ~/jax-metal/bin/activate
python -m pip install -U pip
python -m pip install numpy wheel
python -m pip install jax-metal
Usage
python -c 'import jax; print(jax.numpy.arange(10))'
Use with advanced jaxlib
jax-metal is compatible with the minimal jaxlib version tracked in the above table. It can be compatibly run with advanced jaxlibs no beyond the version with "<", by setting enviroment variable ENABLE_PJRT_COMPATIBILITY. However functionality correctness is not guarenteed under the circumtances.
pip install -U jaxlib jax
ENABLE_PJRT_COMPATIBILITY=1 python -c 'import jax; print(jax.numpy.arange(10))'
Please refer to https://developer.apple.com/metal/jax/ for the full setup and verification.
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
Built Distributions
Hashes for jax_metal-0.0.6-py3-none-macosx_13_0_arm64.whl
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 368a7f5933b7a8278588b1e1e4a5451926176372c4c60d3706a10afa71b8361d |
|
| MD5 | 5f2635e93a897f049ee8c13d35a3d1b1 |
|
| BLAKE2b-256 | a372a24f74df137562e5376f4c5561edba13888e8b88b5b82a37db75abc2004b |
Hashes for jax_metal-0.0.6-py3-none-macosx_10_14_x86_64.whl
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 8b35a9d059e82432a3c5d68643cd243dca3ceea49b218174b79f8d0473f4bf7c |
|
| MD5 | fb0077b5039beca3fbe3d4a3c5e9c1b7 |
|
| BLAKE2b-256 | 8716617a09e900f57aca17b045f590e09ebd188be09c9ef4d28fe999242ee661 |
