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.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 |
For latest jax-metal installation,
python -m pip install jax-metal
For 0.0.3 and before, pls refer to https://developer.apple.com/metal/jax/ for setup.
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.4-py3-none-macosx_11_0_arm64.whl
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 5238c8aa82bd82a1f5c73a3f06e5cb26fe88fb35f7dcdd7ee4d5b1c29a0a5415 |
|
| MD5 | f79e6eb237f8f451249438b040e377cb |
|
| BLAKE2b-256 | 7e59ff91dc65e7f945479b08509185d07de0c947e81c07705367b018cb072ee9 |
Close
Hashes for jax_metal-0.0.4-py3-none-macosx_10_14_x86_64.whl
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 5ac37657c2acf2f35081358ab451618882baaa435d792e77121a0e45f61a4898 |
|
| MD5 | edd36f249bf4760e2171f14ba1f07998 |
|
| BLAKE2b-256 | 7dea05050434826a648d26bdcad4cb02642d548d9803945883414ac522783db3 |
