No worries @dingshuhan, love all the work the Apple team are doing to get JAX to use Apple GPUs! Hope you have a great week

That's a nice hot fix! Thanks for posting mate