PyTorch XLA for TPU kernels

I don’t know if this is the right repository for requesting changes to the TPU Kaggle Kernel VM. If not, please let me know which is the correct GitHub repository or Kaggle forum for this.

When we use PyTorch XLA in Kaggle Kernels, we always require the following installation code:

!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
!python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev

Essentially, this code runs this program. It obtains the appropriate wheels, uninstalls torch and torchvision, and installs torch_xla and its corresponding versions of torch and torchvision. It also installs two Linux packages, libomp5 and libopenblas-dev.

Can the Linux packages and PyTorch XLA package be pre-installed in the TPU Kaggle Kernel? This process takes two minutes or so, and it’s slightly annoying to keep waiting for this during interactive mode. So it would be really helpful if it is pre-installed.

[ifigotin]

@jysohn23 may have more context on this, however, my understanding is that the releases are not GA yet, so may possibly break, and have instabilities. Additionally, the current installation process (via the script provided by the team) will also change. We are planning to revisit integrating everything into Kaggle image after the GA.

[tmabraham]

@ifigotin My understanding though is that there is already a release with most of the major features (unless there is a plan to change a major part of the interface). I am not talking about the nightly version, but rather the v1.5 release from a couple months ago. However, I understand if PyTorch XLA team is not ready for it to be available as part of the Kaggle image, and I guess I will just have to wait…