Setup Nx lib and EXLA to run NX/AXON with CUDA

Steps for you setup and run Machine Learning with Axon or simple Nx script with EXLA on CUDA (GPU).

1. Setup CUDA on local machine.

For case you want to try run model on GPU (Linux/Ubuntu) you need setup CUDA environment follow steps.

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-ubuntu2204-11-8-local_11.8.0-520.61.05-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-11-8-local_118.0-520.61.05-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda

Note: Remember check installed version CUDA is matched to toolkit version by commands

nvcc …

axon case cuda environment etc gpu linux machine machine learning pin setup simple ubuntu