Feature Request: Install the right XLA wheels based on the torch version and environment

With the torch xla wheels now having a separate wheel directly for Google Colab, I think it’d be useful to have an xla installer as part of a utility in Accelerate.

In my head I imagine it would look something like so:

1. Check if we are in Colab or not. (there are ways to do this, see the notebook_launcher code)
2. Grab the torch XLA wheel based on the current installed torch version
3. Install said wheel

An upgrade=True could be passed that will just install the latest torch and torch_xla wheels as well.

API wise it’d look something like so:

# from the notebook
from accelerate.utils import install_xla
install_xla()

# from the CLI
accelerate install_xla

[muellerzr]

Feel free @bipinKrishnan! Yep more or less! If you can get that going and have the pythonic interface I can whip up the CLI half if it’s a bit too confusing 😃

[bipinKrishnan]

@muellerzr Can I work on this issue?

I’m not quite sure about the CLI version but based on what you’ve already mentioned, these are the things to be done for the notebook:

1. Create a separate file for xla stuff in accelerate/utils
2. Add the code for colab check and to get the current version of torch
3. Use python subprocess to install the xla wheels for the current torch version

Let me know if I’m missing something 🙂.