===== JAX =====
Como usar o JAX nas máquinas da Rede IME.

Opção 1: Usar apptainer:
<code bash>
# 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
</code>

A vantagem de usar o apptainer é que o ambiente é bastante controlado e há suporte nativo à GPU Nvidia.

Opção 2: Usar o conda:
<code bash>
# 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]
</code>

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