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

Run a ready-to-use Docker image with JAX 0.11.1 and JupyterLab on top of your ROCm-enabled AMD Instinct system. JAX is installed with AMD’s official ROCm 7.14 PJRT plugin, so jax.jit compiles and executes directly on the GPU. All software is preinstalled in containers, so you can start straight from either SSH or Jupyter.

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.

Getting Started After Deploying JAX

Accessing JAX + ROCm resources and libraries

There are two ways to use the JAX + ROCm stack on this droplet: directly from JupyterLab, or via the Docker container.

• JupyterLab is preferable if you have notebook projects or just want to explore and experiment.
• The Docker container is preferable if you have a Python project that requires deployment.

Option 1: JupyterLab

Open JupyterLab in a browser at http://<your-droplet-ip>.

The access token is printed in the login banner. SSH in as root to read it:

ssh root@<your-droplet-ip>

Option 2: Docker container

JAX runs inside a container named rocm. Open a shell in it:

docker exec -it rocm /bin/bash

Confirm JAX sees the GPU

docker exec -it rocm python -c "import jax; print(jax.__version__); print(jax.devices())"

Expected output on this image:

0.11.1
[RocmDevice(id=0)]

If it reports a CPU device instead, the ROCm plugin did not load.

Run something on the GPU

docker exec -it rocm python -c "import jax, jax.numpy as jnp; x = jnp.ones((4096, 4096)); print(jax.jit(lambda a: (a @ a).sum())(x))"

Inspect the hardware and logs

Check the GPU and ROCm stack from the host with amd-smi or rocminfo. Container logs are available with docker logs rocm.

Networking

The UFW firewall is enabled. Only ports 22 (SSH), 80 (HTTP) and 443 (HTTPS) are open.