For AI agents: The documentation index is at https://docs.digitalocean.com/llms.txt. Markdown versions of pages use the same URL with index.html.md in place of the HTML page (for example, append index.html.md to the directory path instead of opening the HTML document).

JAX 0.8.2 on AMD ROCm 7.2.4

This 1-Click image ships JAX 0.8.2, a high-performance numerical computing and machine learning framework with composable function transformations (autograd, JIT, vmap, pmap), running on an AMD ROCm 7.2.4 host stack on Ubuntu 24.04. It is preconfigured for AMD Instinct GPUs and packaged as a ready-to-run Docker container, alongside a JupyterLab environment with example notebooks.

JAX on ROCm brings XLA-compiled, GPU-accelerated array computing to AMD Instinct hardware for research and production ML workloads.

Software Included

Deploying this Offering using the Control Panel

Click the Deploy to DigitalOcean button to deploy this offering. If you aren’t logged in, this link will prompt you to log in with your DigitalOcean account.

[](https://cloud.digitalocean.com/gpus/new?appId=9e60978f1ebb91bdd44e8e1c&image=JAX 0.8.2 (ROCm 7.2.4) 0.8.2 on Ubuntu 24.04&type=applications)

Getting Started After Deploying JAX 0.8.2 (ROCm 7.2.4)

After the droplet boots, SSH in as root. The Jupyter Lab URL and token, plus the commands to start and inspect the JAX container, are printed in the login MOTD. JAX runs inside a Docker container; use the printed docker exec command to open an interactive shell, then python -c "import jax; print(jax.devices())". Verify the GPU stack with amd-smi and rocminfo.