Metal plug-in JAX uses the new Metal plug-in to provide Metal acceleration on Mac platforms. The Metal plugin uses the OpenXLA compiler and PjRT runtime to accelerate JAX machine learning workloads on GPU. The OpenXLA compiler lowers the JAX primitives to a Stable HLO format, which is converted to MPSGraph executables and Metal runtime APIs to dispatch to GPU. Requirements Mac computers with Apple silicon or AMD GPUs Python 3.9 or later Xcode command-line tools: xcode-select --install Get starte...