Skip to content

Commit

Permalink
Add nightly install instructions to docs
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 576990265
  • Loading branch information
yashk2810 authored and jax authors committed Oct 26, 2023
1 parent 873cffc commit 87c0561
Showing 1 changed file with 31 additions and 0 deletions.
31 changes: 31 additions & 0 deletions docs/installation.md
Expand Up @@ -162,6 +162,37 @@ Toolbox](https://github.com/NVIDIA/JAX-Toolbox) containers, which are
bleeding edge containers containing nightly releases of jax and some
models/frameworks.

## Nightly installation

Nightly releases reflect the state of the main repository at the time they are
built, and may not pass the full test suite.

* JAX:
```bash
pip install -U --pre jax -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
```

* Jaxlib CPU:
```bash
pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
```

* Jaxlib TPU:
```bash
pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
pip install -U libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```

* Jaxlib GPU (Cuda 12):
```bash
pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_cuda12_releases.html
```

* Jaxlib GPU (Cuda 11):
```bash
pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_cuda_releases.html
```

## Google TPU

### pip installation: Google Cloud TPU
Expand Down

0 comments on commit 87c0561

Please sign in to comment.