Accelerated JAX training on Mac with Sonoma 14.1

I have already followed the instruction as follows:

https://developer.apple.com/metal/jax/

to install jax and jaxlib packages in my jupyter notebook, however, it still doesn't work, and I found out that my macOS version is Sonoma 14.1, I wonder whether there are any updates to the instruction?

Thank you!

Accelerated JAX training on Mac with Sonoma 14.1
 
 
Q