Home
ProductsComfyUIMulti-cloud console
PricingBlogsDocs

Does JAX Run on AMD GPUs with ROCm?

Back
Team Aquanode

Team Aquanode

Sarthak Vaish

NOVEMBER 13, 2026

Yes. AMD documents JAX with ROCm support on Linux for its Instinct accelerators, through prebuilt rocm/jax Docker images or a pip install.

What AMD lists

AMD's JAX install guide lists these Instinct GPUs as supported: MI355X, MI350X, MI350P, MI325X, MI300X, MI300A, MI250X, MI250 and MI210. Some Radeon products are listed as well.

Two install paths are described:

  • Docker: pull a rocm/jax image matching your Python and JAX versions, then pass the AMD GPU devices through with --device /dev/kfd --device /dev/dri.
  • pip: install the ROCm libraries and device packages from AMD's package index, then the ROCm JAX plugin and PJRT packages, then jax and jaxlib from PyPI.

AMD notes that "unlike PyTorch, the JAX packages don't automatically install ROCm library and device packages as dependencies." To check the setup, run python -c "import jax; print(jax.devices())" and look for a RocmDevice entry per GPU.

For the PyTorch side, see PyTorch on AMD Instinct. The other alternatives are in NVIDIA's main competitors in AI.

#jax#rocm#amd instinct#amd

Related reading

TPU vs GPU: The Real Architectural and Cost Comparison

TPU vs GPU compared on architecture, current-gen specs, software, and price: Google's Ironwood TPU against NVIDIA's H100/H200, with a clear decision framework.

Is PyTorch Officially Supported on AMD Instinct via ROCm?

PyTorch runs on AMD GPUs through ROCm, with tested wheels and AMD-validated Docker images, but ROCm does not support every PyTorch feature.

AMD MI355X Guide: Specs, MLPerf, vs B200 (2026)

AMD Instinct MI355X specs, MLPerf Inference v6.0 results, MI355X vs MI350X and B200, power and cooling needs, and when it beats an MI300X.

Your own GPU cloud accounts

One console for the GPU clouds you already rent from. Your keys stay on your machine. $20 a month, flat.

Connect your accounts

One console for the GPU clouds you already use.

Connect your accounts for $20 a month, flat. Your provider bills the GPU.

Aquanode LogoAquanode

Run AI without spending any time on infrastructure.

© 2026 Aquanode. All rights reserved.

All trademarks, logos and brand names are the property of their respective owners.