Installing JAX for ROCm

Installing JAX for ROCm#

Applies to Linux

2024-05-21

4 min read time

JAX provides a NumPy-like API, which combines automatic differentiation and the Accelerated Linear Algebra (XLA) compiler to achieve high-performance machine learning at scale.

JAX uses composable transformations of Python+NumPy through just-in-time (JIT) compilation, automatic vectorization, and parallelization.

To learn about JAX, including profiling and optimizations, refer to the JAX documentation.

Compatibility#

You can currently use JAX with the following hardware and software:

  • GPUs: MI250 and MI300

  • OS: Ubuntu 20.04

  • Python: 3.9, 3.10, 3.11

  • ROCm: 6.0.0, 6.1.0

Installing JAX#

JAX wheels and Docker images are released through the GitHub ROCm JAX fork.

Tip

To build JAX from source files, refer to the JAX developer documentation or use the ROCm build script.

  1. Pull the latest public JAX Docker image.

    docker pull rocm/jax:latest
    
  2. Start Docker container.

    docker run -it -w /workspace --device=/dev/kfd --device=/dev/dri --group-add video \
    --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --shm-size 16G rocm/jax:latest
    
  1. Verify the installation.

    python3 -c 'import jax' 2> /dev/null && echo 'Success' || echo 'Failure'
    
  2. Verify that the GPU is accessible from JAX.

    python3 -c 'import jax; print(jax.devices())'
    
  3. Run a basic example to ensure installation is successful.

    git clone https://github.com/google/jax.git jax
    cd jax
    export PYTHONPATH=/workspace/jax/examples/:$PYTHONPATH
    python3 examples/mnist_classifier.py
    

    Your output should look similar to this:

    Starting training...
    Epoch 0 in 10.97 sec
    Training set accuracy 0.871916651725769
    Test set accuracy 0.880299985408783
    Epoch 1 in 0.34 sec
    Training set accuracy 0.8979166746139526
    Test set accuracy 0.9030999541282654
    Epoch 2 in 0.33 sec
    Training set accuracy 0.9092333316802979
    Test set accuracy 0.9142999649047852
    Epoch 3 in 0.33 sec
    Training set accuracy 0.9170833230018616
    Test set accuracy 0.9220999479293823
    Epoch 4 in 0.33 sec
    Training set accuracy 0.9226333498954773
    Test set accuracy 0.9279999732971191
    Epoch 5 in 0.33 sec
    Training set accuracy 0.9271667003631592
    Test set accuracy 0.9297999739646912
    Epoch 6 in 0.34 sec
    Training set accuracy 0.9323500394821167
    Test set accuracy 0.9328999519348145
    Epoch 7 in 0.34 sec
    Training set accuracy 0.935699999332428
    Test set accuracy 0.9364999532699585
    Epoch 8 in 0.33 sec
    Training set accuracy 0.938800036907196
    Test set accuracy 0.9393999576568604
    Epoch 9 in 0.33 sec
    Training set accuracy 0.9425833225250244
    Test set accuracy 0.9418999552726746