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.5 | Sonoma 14.2+ | v0.4.20 | 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
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
No source distribution files available for this release.See tutorial on generating distribution archives.
Built Distributions
Close
Hashes for jax_metal-0.0.5-py3-none-macosx_11_0_arm64.whl
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 565e211628059ca44021349ec36e514bcbfcbf33af610b2912829569bf78b596 |
|
| MD5 | 7a4249ea5c248b9c224f94d3538bcd5c |
|
| BLAKE2b-256 | 8e15e0c6c864b5a8210d7c0bec34a567b3b8d6071422412333b7fff24c324bea |
Close
Hashes for jax_metal-0.0.5-py3-none-macosx_10_14_x86_64.whl
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 08059fe8d12897150bf49798b7cded9e1afe9953a9fd710a32e2713f301fee8c |
|
| MD5 | 080da2901cab0fd70192637d66c4c694 |
|
| BLAKE2b-256 | 3905213194270f3836fe81f6220f5f23a9dc3132cb3a832547af529059fb4514 |
