Mac 上加速 JAX：Metal

本文介绍了如何在Mac平台上使用Metal插件加速JAX机器学习工作负载。Metal插件利用OpenXLA编译器和PjRT运行时，将JAX原语转换为Stable HLO格式，再转换为MPSGraph可执行文件和Metal运行时API，最终在GPU上运行。文章列出了jax-metal、macOS和jaxlib的版本兼容性，并提供了安装和验证步骤，以及使用AXLearn Fuji模型进行推理的示例。文章还指出，Metal插件仍处于实验阶段，部分JAX功能和数据类型尚不支持。