From ace5e24b96f0012a311a724676fdaaa8f0f86d1c Mon Sep 17 00:00:00 2001 From: Siyuan Liu Date: Fri, 7 Feb 2025 13:58:12 -0800 Subject: [PATCH] [torchax] Update installation guide in readme (#8664) Co-authored-by: Siyuan Liu --- torchax/README.md | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/torchax/README.md b/torchax/README.md index ff38ea5d75bb..797e1f3f8553 100644 --- a/torchax/README.md +++ b/torchax/README.md @@ -26,7 +26,7 @@ The following instructions assume you are in the `torchax` directory: ``` Fork the repository $ git clone https://github.com//xla.git -$ cd xla/experimental/torchax +$ cd xla/torchax ``` @@ -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