# 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.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))'
```
### 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.
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## 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.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### Use with advanced jaxlib\njax-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.\n```\npip install -U jaxlib jax\nENABLE_PJRT_COMPATIBILITY=1 python -c 'import jax; print(jax.numpy.arange(10))'\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.0.7",
"project_urls": {
"Homepage": "https://developer.apple.com/metal/jax/"
},
"split_keywords": [
"jax",
"metal",
"plugin"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "fbb08ad3c3a71351dc7e6bf98d15058c81890a04544835aa2b986150bb8f01d0",
"md5": "7b028e9861246b81f9020d3c5c97fb15",
"sha256": "9b1b90221aa5718b88b9d0c97449fc95a8594b9b609bde086206a6742510831e"
},
"downloads": -1,
"filename": "jax_metal-0.0.7-py3-none-macosx_10_14_x86_64.whl",
"has_sig": false,
"md5_digest": "7b028e9861246b81f9020d3c5c97fb15",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": null,
"size": 55023918,
"upload_time": "2024-05-04T00:36:59",
"upload_time_iso_8601": "2024-05-04T00:36:59.457718Z",
"url": "https://files.pythonhosted.org/packages/fb/b0/8ad3c3a71351dc7e6bf98d15058c81890a04544835aa2b986150bb8f01d0/jax_metal-0.0.7-py3-none-macosx_10_14_x86_64.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "d64ff5d128a493b7387fbbe0e6906544214af2a6b86af30302dd6ffb9dc66a74",
"md5": "cd8a3eab914ed269b8717554293f8c76",
"sha256": "0f356de48d721ecc8a697ed97ae88b9c1323e546614a5b82ddb4721d4deb1021"
},
"downloads": -1,
"filename": "jax_metal-0.0.7-py3-none-macosx_13_0_arm64.whl",
"has_sig": false,
"md5_digest": "cd8a3eab914ed269b8717554293f8c76",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": null,
"size": 40797644,
"upload_time": "2024-05-04T00:35:43",
"upload_time_iso_8601": "2024-05-04T00:35:43.941607Z",
"url": "https://files.pythonhosted.org/packages/d6/4f/f5d128a493b7387fbbe0e6906544214af2a6b86af30302dd6ffb9dc66a74/jax_metal-0.0.7-py3-none-macosx_13_0_arm64.whl",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-05-04 00:36:59",
"github": false,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"lcname": "jax-metal"
}
JSON