---
title: JAX 0.8.2 (ROCm 7.2.4)
description: JAX 0.8.2 high-performance numerical computing and ML framework on AMD ROCm 7.2.4 host, Ubuntu 24.04, ready for AMD Instinct accelerators.
product: Marketplace
url: https://docs.digitalocean.com/products/marketplace/catalog/jax-0-8-2-rocm-7-2-4/
last_updated: "2026-06-30"
---

> **For AI agents:** The documentation index is at [https://docs.digitalocean.com/llms.txt](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 (ROCm 7.2.4)

Generated on 30 Jun 2026 from [the JAX 0.8.2 (ROCm 7.2.4) catalog page](https://marketplace.digitalocean.com/apps/jax-0-8-2-rocm-7-2-4)

## 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

| Package | Version | License |
|---|---|---|
| [JAX](https://docs.jax.dev) | 0.8.2 | Apache-2.0 |
| [ROCm](https://rocm.docs.amd.com) | 7.2.4 | MIT/Apache-2.0 |
| [JupyterLab](https://jupyter.org) | 4.4.2 | |
| Docker | latest | |

## 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.

\[![Deploy JAX 0.8.2 (ROCm 7.2.4) to DO](https://www.deploytodo.com/do-btn-blue.svg)]([https://cloud.digitalocean.com/gpus/new?appId=9e60978f1ebb91bdd44e8e1c&image=JAX](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`.