Programmatically find jax backend

[Ruturaj4]

Hello,

what is the best way to find the JAX backend. Whether it is rocm or cuda for e.g.

We used to have something like jaxlib.rocm before. But I think that has changed in recent versions.

Is there an alternative method (maybe pjrt query) to verify if the JAX/JAXlib installation supports ROCm during the Docker image build process? Checking with jax.devices() works once the container is running, but it isn’t useful during the image build stage.