{"library":"torchax","install":[{"cmd":"pip install torchax","imports":["import torchax","torchax.enable_globally()","from torchax.interop import JittableModule","from torchax.interop import jax_jit"]},{"cmd":"# First, install PyTorch CPU:\npip install torch --index-url https://download.pytorch.org/whl/cpu # Linux\npip install torch # Mac\n\n# Then, install JAX for your accelerator:\npip install -U jax[tpu] # Google Cloud TPU\npip install -U jax[cuda12] # GPU machines\npip install -U jax # Linux CPU or Mac","imports":[]}]}