I believe all big frameworks have some work being done to make them compatible with AMD GPUs. Here is the relevant issue for JAX (support seem to be in alpha bu

I believe all big frameworks have some work being done to make them compatible with AMD GPUs. Here is the relevant issue for JAX (support seem to be in alpha but viable): https://github.com/google/jax/issues/2012 https://github.com/google/jax/issues/2012