rocm/jax repository overview

⁠JAX-ROCm Docker Images:

This repository contains the latest support for JAX in ROCm framework.

Here's the recommended command to launch JAX-ROCm docker image:

alias drun='sudo docker run -it --network=host --device=/dev/kfd --device=/dev/dri --ipc=host --shm-size 16G --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -w /root'

drun rocm/jax:latest