JAX has matured a lot over last 6 years and XLA has been around for lot longer. We believe we can extract good perf from AMD with JAX + XLA kernels.