Skip to content

Commit

Permalink
[torchax] Update installation guide in readme (#8664)
Browse files Browse the repository at this point in the history
Co-authored-by: Siyuan Liu <[email protected]>
  • Loading branch information
lsy323 and Siyuan Liu authored Feb 7, 2025
1 parent 0cd1fc2 commit ace5e24
Showing 1 changed file with 3 additions and 9 deletions.
12 changes: 3 additions & 9 deletions torchax/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ The following instructions assume you are in the `torchax` directory:
```
Fork the repository
$ git clone https://github.com/<github_username>/xla.git
$ cd xla/experimental/torchax
$ cd xla/torchax
```


Expand Down Expand Up @@ -55,17 +55,11 @@ Note: `dev-requirements.txt` will install the CPU-only version of PyTorch.

#### 1.1 Install this package

If you want to install torchax without the jax dependency and use the jax dependency from torch_xla:
```bash
pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
pip install -e .
```

Otherwise, install `torchax` from source for your platform:
Install `torchax` from source for your platform:
```bash
pip install -e .[cpu]
pip install -e .[cuda]
pip install -e .[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html
pip install -e .[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```

#### 1.2 (optional) verify installation by running tests
Expand Down

0 comments on commit ace5e24

Please sign in to comment.