JAX加速Mac GPU。
项目描述
jax-metal
jax-metal软件包是一个Metal GPU插件，用于为JAX应用程序在Mac平台上提供Metal加速。
安装
以下表格跟踪jax-metal版本以及与jax、jaxlib和MacOS兼容的版本。
| jax-metal | MacOS | jaxlib | jax |
|---|---|---|---|
| 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 测试版 | >=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 测试版 | v0.4.10 | v0.4.11 |
我们建议使用venv或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
用法
python -c 'import jax; print(jax.numpy.arange(10))'
与jaxlib的兼容性
jax-metal 兼容上表中跟踪的最小 jaxlib 版本。可以通过设置环境变量 ENABLE_PJRT_COMPATIBILITY=1 与高于最低版本的 jaxlibs 兼容运行。
pip install -U jaxlib jax
ENABLE_PJRT_COMPATIBILITY=1 python -c 'import jax; print(jax.numpy.arange(10))'
目前不支持
Metal 插件是实验性的，可能不支持所有 JAX 功能。报告并跟踪的问题可以在以下列表中找到：https://github.com/google/jax/issues?q=is%3Aissue+is%3Aopen+metal
- 不支持的数据类型：np.float64、np.complex64、np.complex128
- Metal 插件在https://github.com/google/jax/tree/main/tests 下的所有测试都无法通过。
有关完整设置和验证，请参阅https://developer.apple.com/metal/jax/。
关闭
散列 for jax_metal-0.1.0-py3-none-macosx_11_0_arm64.whl
| 算法 | 散列摘要 | |
|---|---|---|
| SHA256 | 68006a39d331674880265361f612fcb0e3314a220245c05b8c27d204077de085 |
|
| MD5 | b453eb8d3cc560c0f12a80aa8c85ef8f |
|
| BLAKE2b-256 | 80afed482a421a868726e7ca3f51ac19b0c9a8e37f33f54413312c37e9056acc |
关闭
散列 for jax_metal-0.1.0-py3-none-macosx_10_14_x86_64.whl
| 算法 | 散列摘要 | |
|---|---|---|
| SHA256 | 6cbd62a2068fcadea9c22689c12905e6fe0f58da2230fe2beba263132258cc54 |
|
| MD5 | 88c679698cfe2386efb76d9769e07299 |
|
| BLAKE2b-256 | 516a1c0e2d07d92c6583e874ef2bbf4382662a3469bbb661d885eeaaddca426f |
