The description says it's autograd + XLA. So I assumed it always compiles to GPUs via XLA. But, had a look in the code and jax has cublas and RoCm blas, and it

The description says it's autograd + XLA. So I assumed it always compiles to GPUs via XLA. But, had a look in the code and jax has cublas and RoCm blas, and it looks like there is a flow where it uses the gpu directly, unless I'm missing something. Definitely worth having a closer look. Autograd via function reflection should be faster than backprop. And if it's running on AMD GPUs then it's quite intriguing.