Como usar o JAX nas máquinas da Rede IME.
Opção 1: Usar apptainer:
# No terminal, rodar os comandos: apptainer pull docker://ghcr.io/nvidia/jax:jax # Será criado um arquivo chamado jax_jax.sif apptainer run --nv jax_jax.sif
A vantagem de usar o apptainer é que o ambiente é bastante controlado e há suporte nativo à GPU Nvidia.
Opção 2: Usar o conda:
# criar o ambiente conda create -n jax conda init bash source .bashrc conda activate jax conda install pip # instalar o jax pip install jax[cpu] # se houver GPU da Nvidia no computador usar jax[cuda12] em vez de jax[cpu]
A vantagem de usar o conda é usar o procedimento de instalação descrito na documentação oficial.