XLA Check Failed: options.is_autotuning_compilation about jax

[sjrothfuss]

We solved this issue by upgrading CUDA on our machine.

NVIDIA-SMI 550.54.14 Driver Version: 550.54.14 CUDA Version: 12.4

If someone else has the same issue and can't upgrade, I recommend specifying cuDNN and CUDA Toolkit versions before installing, such as by running the following.

conda update -y -n base -c defaults conda
conda create -n JAXENV -y python=3.8
conda activate JAXENV
conda install -y -c nvidia cudnn=8.0.4
conda install -y -c conda-forge cudatoolkit-dev=11.1.1
chmod 755 /miniconda3/envs/JAXENV/pkgs/cuda-toolkit
pip install --upgrade pip
pip install --upgrade jax==0.3.25 jaxlib==0.3.25+cuda11.cudnn805 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

(This code installs CUDA 11 and the compatible JAX because conda-forge does not have toolkit version ≥ 12 at the time of writing.)