Wiki da Rede IME

JAX

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.

https://jax.readthedocs.io/en/latest/installation.html