Jax GPU access for us Mac users with the new M1/M2 silicon has been non-existent. This might change with the recent release by Apple of jax-metal . I have begun trying to get it to work but wanted to post this here in case other people wanted to try it as well. If I get it working, I will post my s…

Jax on Apple Silicon GPU

zweli June 13, 2023, 6:24am 2

I got the jax-metal library to see my M1 GPU

Screenshot 2023-06-13 at 2.22.18 AM

1 Like