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
File details
Details for the file jax_metal-0.0.4-py3-none-macosx_11_0_arm64.whl.
File metadata
- Download URL: jax_metal-0.0.4-py3-none-macosx_11_0_arm64.whl
- Upload date:
- Size: 38.7 MB
- Tags: Python 3, macOS 11.0+ ARM64
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.2 CPython/3.11.4
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 5238c8aa82bd82a1f5c73a3f06e5cb26fe88fb35f7dcdd7ee4d5b1c29a0a5415 |
|
| MD5 | f79e6eb237f8f451249438b040e377cb |
|
| BLAKE2b-256 | 7e59ff91dc65e7f945479b08509185d07de0c947e81c07705367b018cb072ee9 |
Provenance
File details
Details for the file jax_metal-0.0.4-py3-none-macosx_10_14_x86_64.whl.
File metadata
- Download URL: jax_metal-0.0.4-py3-none-macosx_10_14_x86_64.whl
- Upload date:
- Size: 51.9 MB
- Tags: Python 3, macOS 10.14+ x86-64
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.2 CPython/3.11.4
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 5ac37657c2acf2f35081358ab451618882baaa435d792e77121a0e45f61a4898 |
|
| MD5 | edd36f249bf4760e2171f14ba1f07998 |
|
| BLAKE2b-256 | 7dea05050434826a648d26bdcad4cb02642d548d9803945883414ac522783db3 |
