I'm glad to see a full implementation on AMD hardware. I'm not familiar with JAX, but the idea of providing an abstraction layer to more easily get to work on

I'm glad to see a full implementation on AMD hardware. I'm not familiar with JAX, but the idea of providing an abstraction layer to more easily get to work on what hardware is available seems really valuable. Bringing back some competitiveness to the ecosystem will be a big win for workload mobility. I suspect that price/performance across implementations will be highly dependent on contract details, but do you intend to publish some comparisons in the future?