{"library":"jaxlib","install":[{"cmd":"pip install --upgrade pip\npip install --upgrade jax jaxlib","imports":["import jax","import jax.numpy as jnp"]},{"cmd":"pip install --upgrade pip\npip install --upgrade \"jax[cuda12_pip]\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html","imports":[]},{"cmd":"# For specific CUDA versions or locally installed CUDA, refer to official JAX docs:\n# https://jax.readthedocs.io/en/latest/installation.html","imports":[]}]}