Repository: huggingface/accelerate

Issue:
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:

```python
# from the notebook
from accelerate.utils import install_xla
install_xla()
```

```bash
# from the CLI
accelerate install_xla
```




Discussion:
Hi, can I work on this? :)
Yes you can @Atharva7K! :D 
@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 :slightly_smiling_face:.
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 :)

Resolving patch:
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -128,4 +128,5 @@
     write_basic_config,
 )
 from .random import set_seed, synchronize_rng_state, synchronize_rng_states
+from .torch_xla import install_xla
 from .tqdm import tqdm
diff --git a/src/accelerate/utils/torch_xla.py b/src/accelerate/utils/torch_xla.py
new file mode 100644
--- /dev/null
+++ b/src/accelerate/utils/torch_xla.py
@@ -0,0 +1,44 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import subprocess
+import sys
+
+import pkg_resources
+
+
+def install_xla(upgrade: bool = False):
+    """
+    Helper function to install appropriate xla wheels based on the `torch` version.
+
+    Args:
+        upgrade (`bool`, *optional*, defaults to `False`):
+            Whether to upgrade `torch` and install the latest `torch_xla` wheels.
+    """
+    in_colab = False
+    if "IPython" in sys.modules:
+        in_colab = "google.colab" in str(sys.modules["IPython"].get_ipython())
+
+    if in_colab:
+        if upgrade:
+            torch_install_cmd = ["pip", "install", "-U", "torch"]
+            subprocess.run(torch_install_cmd, check=True)
+        # get the current version of torch
+        torch_version = pkg_resources.get_distribution("torch").version
+        torch_version_trunc = torch_version[: torch_version.rindex(".")]
+        xla_wheel = f"https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-{torch_version_trunc}-cp37-cp37m-linux_x86_64.whl"
+        xla_install_cmd = ["pip", "install", xla_wheel]
+        subprocess.run(xla_install_cmd, check=True)
+    else:
+        raise RuntimeError("`install_xla` utility works only on google colab.")


Test patch: