#!/bin/bash

# This test checks that JAX
# can be used to run an example

# NOTE: this test assumes JAX has been installed according
# to the instructions available in the model installation repo:
# https://github.com/amd/HPCTrainingDock/blob/main/extras/scripts/jax_setup.sh


module -t list 2>&1 | grep -q "^rocm"
if [ $? -eq 1 ]; then
  echo "rocm module is not loaded"
  echo "loading default rocm module"
  module load rocm
fi
module load jax

SRC_DIR=$(pwd)
BUILD_DIR=$(mktemp -d)
trap "rm -rf ${BUILD_DIR}" EXIT
cp * ${BUILD_DIR}

cd ${BUILD_DIR}

git clone https://github.com/google/jax.git jax
cp jax/examples/datasets.py .
cp jax/examples/mnist_classifier.py .
sed -i -e 's/from examples //' mnist_classifier.py
export PYTHONPATH=.:$PYTHONPATH
python3 mnist_classifier.py
