We appreciate all contributions. If you are planning to contribute a bug fix for an open issue, please comment on the thread and we're happy to provide any guidance. You are very welcome to pick issues from good first issue and help wanted labels.
If you plan to contribute new features, utility functions or extensions to the core, please first open an issue and discuss the feature with us. Sending a PR without discussion might end up resulting in a rejected PR, because we might be taking the core in a different direction than you might be aware of.
We recommend you to use our prebuilt Docker image to start your development work using one of the two following methods.
-
Create an empty directory (optionally on a remote host via SSH) and open it in VSCode. Then, clone PyTorch, TorchVision, and PyTorch/XLA:
git clone --recursive --depth=1 https://github.com/pytorch/pytorch.git # Optional: install TorchVision if you need to run tests that involve vision modules git clone --recursive --depth=1 https://github.com/pytorch/vision.git git clone https://github.com/pytorch/xla.git pytorch/xla # Optional: use git@github.com:pytorch/xla.git instead if you prefer to use SSH with key forwarding
-
Link (or copy) VSCode configuration to your workspace directory:
ln -s pytorch/xla/.devcontainer/ .devcontainer ln -s pytorch/xla/contrib/vscode/ .vscode ln -s pytorch/xla/.style.yapf .style.yapf ln -s pytorch/xla/.clang-format .clang-format
-
From VSCode's command menu, run
Reopen in Container
from the command palette (F1 key) to open your workspace in one of our pre-built Docker containers. Select the correct container config based on your local accelerator (default totpu-contributor
if you are not sure).- If you cannot find
Reopen in Container
, make sure theDev Containers
VSCode extension is installed, then open thepytorch/xla
folder in VSCode.
- If you cannot find
-
Since you are running as root in this container, teach
git
to recognize the repositories you just cloned (outside of docker) as safe:git config --global --add safe.directory /workspaces/torch/pytorch git config --global --add safe.directory /workspaces/torch/pytorch/xla git config --global --add safe.directory /workspaces/torch/vision
-
Build PyTorch, TorchVision, and PyTorch/XLA:
cd pytorch # pytorch/xla requires pytorch wheel to be presented under pytorch/dist python setup.py bdist_wheel python setup.py install cd .. cd vision python setup.py develop cd .. cd pytorch/xla python setup.py develop # Optional: if you're using TPU, install libtpu pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html
-
Test your build
python -c 'import torch_xla as xla; print(xla.device())' # Output: xla:0
-
Setup Development Docker Image
docker pull us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/development:tpu docker run --privileged --name ptxla -it -d -e "TERM=xterm-256color" us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/development:tpu docker exec --privileged -it ptxla /bin/bash
All of the code below will be assumed to be run within the docker.
-
Clone the PyTorch repo as per instructions.
git clone --recursive https://github.com/pytorch/pytorch cd pytorch/
-
Clone the PyTorch/XLA repo:
git clone --recursive https://github.com/pytorch/xla.git
-
Build PyTorch
# pytorch/xla requires pytorch wheel to be presented under pytorch/dist python setup.py bdist_wheel python setup.py develop
-
Build PyTorch/XLA
cd xla/ python setup.py develop
Please refer to this guide.
In pytorch/xla
repo we enforce coding style for both C++ and Python files. Please try to format
your code before submitting a pull request.
pytorch/xla
uses clang-format-11
with a customized style config.
If your PR touches the C++ source files, please run the following command before submitting a PR.
# How to install: sudo apt install clang-format-11
# If your PR only changes foo.cpp, run the following in xla/ folder
clang-format-11 -i -style=file /PATH/TO/foo.cpp
# To format all cpp files, run the following in xla/ folder
find -name '*.cpp' -o -name '*.h' -o -name '*.cc' | xargs clang-format-11 -i -style=file
pytorch/xla
uses yapf
(specially version 0.30.0 in case it's not backward compatible) with a customized style config.
If your PR touches the Python source files, please run the following command before submitting a PR.
# How to install: pip install yapf==0.30.0
yapf --recursive -i *.py test/ scripts/ torch_xla/ benchmarks/
To run the tests, follow one of the options below:
-
Run on local CPU:
export PJRT_DEVICE=CPU
-
Run on Cloud TPU:
export PJRT_DEVICE=TPU
-
Run on GPU:
export PJRT_DEVICE=CUDA GPU_NUM_DEVICES=${NUM_GPU}
For more detail on configuring the runtime, please refer to this doc
If you are planning to be building from source and hence using the latest PyTorch/TPU code base, it is suggested for you to select the Nightly builds when you create a Cloud TPU instance.
Then run test/run_tests.sh
and test/cpp/run_tests.sh
to verify the setup is working.
- If local changes aren't visible, uninstall existing pytorch/xla with
pip uninstall torch_xla
andpip uninstall torch
, then rebuild PyTorch and PyTorch/XLA withpython setup.py develop
orpython setup.py install
. - PJRT errors when running on TPU such as
The PJRT plugin has PJRT API version 0.34. The framework PJRT API version is 0.40
. You need to update yourlibtpu.so
and ensure it's in yourLD_LIBRARY_PATH
environmental directory. You can download a newlibtpu.so
at Google Cloud, which are sorted by date. Download the newest one and install it atpip install libtpu...whl
.