# jax-metal
The jax-metal package is a Metal GPU plugin to provide Metal acceleration on Mac platforms for JAX applications.
## Release Notes
### v0.1.1
The patch fixes IR ops regarding to TopKOp, TanOp, ErfOp, which have changed since jaxlib-0.4.30.
## Installation
The following table tracks `jax-metal` versions and compatible versions of `jax`, `jaxlib` and `MacOS`.
| jax-metal | MacOS | jaxlib | jax |
|-----------|-----------------------------------|---------------------|-----------|
| 0.1.1 | Sonoma 14.4+ | >=0.4.34 | ==jaxlib |
| 0.1.0 | Sonoma 14.4+ | >=0.4.26 | >=0.4.26 |
| 0.0.7 | Sonoma 14.4+ | >=0.4.26 | >=0.4.26 |
| 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))'
```
### Compatibility with jaxlib
jax-metal is compatible with the minimal jaxlib version tracked in the above table. It can be compatibly run with jaxlibs beyond the minimum version by setting the environment variable to ENABLE_PJRT_COMPATIBILITY=1.
```
pip install -U jaxlib jax
ENABLE_PJRT_COMPATIBILITY=1 python -c 'import jax; print(jax.numpy.arange(10))'
```
### Currently not supported
The Metal plug-in is experimental and not all JAX functionality may be supported. Issues that are reported and tracked can be found in the list: https://github.com/google/jax/issues?q=is%3Aissue+is%3Aopen+metal
* Unsupported data types: np.float64, np.complex64, np.complex128
* The Metal plug-in doesn’t pass all tests under https://github.com/google/jax/tree/main/tests.
Please refer to https://developer.apple.com/metal/jax/ for the full setup and verification.
Raw data
{
"_id": null,
"home_page": "https://developer.apple.com/metal/jax/",
"name": "jax-metal",
"maintainer": null,
"docs_url": null,
"requires_python": null,
"maintainer_email": null,
"keywords": "jax metal plugin",
"author": null,
"author_email": null,
"download_url": null,
"platform": null,
"description": "# jax-metal\n\nThe jax-metal package is a Metal GPU plugin to provide Metal acceleration on Mac platforms for JAX applications.\n\n## Release Notes\n\n### v0.1.1\nThe patch fixes IR ops regarding to TopKOp, TanOp, ErfOp, which have changed since jaxlib-0.4.30.\n\n## Installation\n\nThe following table tracks `jax-metal` versions and compatible versions of `jax`, `jaxlib` and `MacOS`.\n\n| jax-metal | MacOS | jaxlib | jax |\n|-----------|-----------------------------------|---------------------|-----------|\n| 0.1.1 | Sonoma 14.4+ | >=0.4.34 | ==jaxlib |\n| 0.1.0 | Sonoma 14.4+ | >=0.4.26 | >=0.4.26 |\n| 0.0.7 | Sonoma 14.4+ | >=0.4.26 | >=0.4.26 |\n| 0.0.6 | Sonoma 14.4 Beta | >=v0.4.22, <v0.4.24 | >=v0.4.22 |\n| 0.0.5 | Sonoma 14.2+ | >=v0.4.20, <v0.4.22 | >=v0.4.20 |\n| 0.0.4 | Sonoma 14.0+ | v0.4.11 | v0.4.11 |\n| 0.0.3 | Ventura 13.4.1+, Sonoma 14.0 Beta | v0.4.10 | v0.4.11 |\n\nWe recommend to install the binary package with venv or conda.\n```\npython3 -m venv ~/jax-metal\nsource ~/jax-metal/bin/activate\npython -m pip install -U pip\npython -m pip install numpy wheel\npython -m pip install jax-metal\n``` \n\n## Usage\n```\npython -c 'import jax; print(jax.numpy.arange(10))'\n```\n### Compatibility with jaxlib\njax-metal is compatible with the minimal jaxlib version tracked in the above table. It can be compatibly run with jaxlibs beyond the minimum version by setting the environment variable to ENABLE_PJRT_COMPATIBILITY=1.\n```\npip install -U jaxlib jax\nENABLE_PJRT_COMPATIBILITY=1 python -c 'import jax; print(jax.numpy.arange(10))'\n```\n### Currently not supported\nThe Metal plug-in is experimental and not all JAX functionality may be supported. Issues that are reported and tracked can be found in the list: https://github.com/google/jax/issues?q=is%3Aissue+is%3Aopen+metal\n\n* Unsupported data types: np.float64, np.complex64, np.complex128 \n* The Metal plug-in doesn\u2019t pass all tests under https://github.com/google/jax/tree/main/tests. \n\n\nPlease refer to https://developer.apple.com/metal/jax/ for the full setup and verification. \n\n",
"bugtrack_url": null,
"license": "Copyright \u00a9 2023 Apple Inc. All rights reserved.",
"summary": "JAX acceleration for Mac GPUs.",
"version": "0.1.1",
"project_urls": {
"Homepage": "https://developer.apple.com/metal/jax/"
},
"split_keywords": [
"jax",
"metal",
"plugin"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "87ec9bb7f7f0ffd06c3fb89813126b2f698636ac7a4263ed7bdd1ff7d7c94f8f",
"md5": "81c3cbc8f16c217ad8b125a990941214",
"sha256": "d918a78443cb808c9491a24a5c2a94cc4eabfd0461d5bcda29a8f332dfbe9b7e"
},
"downloads": -1,
"filename": "jax_metal-0.1.1-py3-none-macosx_10_14_x86_64.whl",
"has_sig": false,
"md5_digest": "81c3cbc8f16c217ad8b125a990941214",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": null,
"size": 54662235,
"upload_time": "2024-10-08T16:57:30",
"upload_time_iso_8601": "2024-10-08T16:57:30.128499Z",
"url": "https://files.pythonhosted.org/packages/87/ec/9bb7f7f0ffd06c3fb89813126b2f698636ac7a4263ed7bdd1ff7d7c94f8f/jax_metal-0.1.1-py3-none-macosx_10_14_x86_64.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "09dc6d8fbfc29d902251cf333414cf7dcfaf4b252a9920c881354584ed36270d",
"md5": "3435ef5a9a7f2a01218264c67c61ea5d",
"sha256": "f1dbfecb298cdd3ba6da3ad6dc9a2adb63d71741f8b8ece28c296b32d608b6c8"
},
"downloads": -1,
"filename": "jax_metal-0.1.1-py3-none-macosx_13_0_arm64.whl",
"has_sig": false,
"md5_digest": "3435ef5a9a7f2a01218264c67c61ea5d",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": null,
"size": 41179678,
"upload_time": "2024-10-08T16:56:31",
"upload_time_iso_8601": "2024-10-08T16:56:31.563808Z",
"url": "https://files.pythonhosted.org/packages/09/dc/6d8fbfc29d902251cf333414cf7dcfaf4b252a9920c881354584ed36270d/jax_metal-0.1.1-py3-none-macosx_13_0_arm64.whl",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-10-08 16:57:30",
"github": false,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"lcname": "jax-metal"
}
JSON