Skip to content

Commit

Permalink
update torch-xla related in READMEs
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 706015739
  • Loading branch information
chunnienc authored and copybara-github committed Dec 13, 2024
1 parent 50c43a3 commit dbd40c7
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 23 deletions.
10 changes: 0 additions & 10 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -102,16 +102,6 @@ Alternately, the nightly version can be installed with:
pip install ai-edge-torch-nightly
```

### Update LD_LIBRARY_PATH if necessary

Torch XLA builds a shared library, `_XLAC.so` that needs to link to the version of Python
it was built with (currently 3.10 or 3.11). In order to ensure that `import _XLAC` can succeed,
update the LD_LIBRARY_PATH to the lib directory of your Python environment:

```bash
export LD_LIBRARY_PATH=<path to Python installation>/lib:$LD_LIBRARY_PATH
```


* The list of versioned releases can be seen [here](https://github.com/google-ai-edge/ai-edge-torch/releases).
* The full list of PyPi releases (including nightly builds) can be seen [here](https://pypi.org/project/ai-edge-torch/#history).
Expand Down
33 changes: 20 additions & 13 deletions docs/pytorch_converter/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -327,19 +327,26 @@ pip install ai-edge-model-explorer
model-explorer 'resnet.tflite'
```

# Use ODML Torch Conversion Backend (Experimental)
# Use Torch XLA Conversion Backend (Legacy)

ODML Torch is an experimental conversion backend for AI Edge Torch, which enables
conversion on macOS and generates TFLite Flatbuffer with higher quality and better
performance. To use AI Edge Torch with ODML Torch, setup your environment these
commands:
AI Edge Torch has been switched to a modern conversion backend for better
on-device compatibility and performance. However, if you encounter compatibility
issues, you can optionally switch to the legacy Torch XLA backend:

```bash
# Install ODML Torch backend dependencies
pip install -r https://raw.githubusercontent.com/google-ai-edge/ai-edge-torch/main/odmltorch-requirements.txt
# Install ai-edge-torch with no default backend (Torch XLA) dependencies
pip install ai-edge-torch-nightly --no-deps

# Disable default AI Edge Torch backend (Torch XLA)
export USE_TORCH_XLA=0
```
# Install ai-edge-torch with torch-xla dependency
pip install ai-edge-torch-nightly[torch-xla]

# Enable torch-xla as the AI Edge Torch backend
export USE_TORCH_XLA=1
```

## Update LD_LIBRARY_PATH if necessary

Torch XLA builds a shared library, `_XLAC.so` that needs to link to the version of Python
it was built with (currently 3.10 or 3.11). In order to ensure that `import _XLAC` can succeed,
update the LD_LIBRARY_PATH to the lib directory of your Python environment:

```bash
export LD_LIBRARY_PATH=<path to Python installation>/lib:$LD_LIBRARY_PATH
```

0 comments on commit dbd40c7

Please sign in to comment.