Migrate uses of import torch_xla as xla to import torch_xla (#9325)

Author: ghpvnist

CONTRIBUTING.md

@@ -166,7 +166,7 @@ commands on your Linux machine directly, outside of the container.
     installed correctly:
 
     ```bash
-    python -c 'import torch_xla as xla; print(xla.device())'
+    python -c 'import torch_xla; print(torch_xla.device())'
     # Output: xla:0
     ```
 
@@ -375,11 +375,11 @@ First, for the `pytorch` repo:
 cd $WORKSPACE_DIR/pytorch
 # Fetch the latest changes from upstream.
 git fetch upstream
-git checkout main 
+git checkout main
 # Merge the changes from upstream/main into your local branch.
 git merge upstream/main
 # Update submodules to match the latest changes.
-git submodule update --recursive 
+git submodule update --recursive
 # Push the updated branch to your fork on GitHub.
 git push origin main
 ```
@@ -389,7 +389,7 @@ Next, for the `vision` repo:
 ```bash
 cd $WORKSPACE_DIR/vision
 git fetch upstream
-git checkout main 
+git checkout main
 git merge upstream/main
 git push origin main
 ```

README.md

@@ -184,7 +184,7 @@ If you're using `DistributedDataParallel`, make the following changes:
 ```diff
  import torch.distributed as dist
 -import torch.multiprocessing as mp
-+import torch_xla as xla
++import torch_xla
 +import torch_xla.distributed.xla_backend
 
  def _mp_fn(rank):
@@ -203,8 +203,8 @@ If you're using `DistributedDataParallel`, make the following changes:
 -  ddp_model = DDP(model, device_ids=[rank])
 
    for inputs, labels in train_loader:
-+    with xla.step():
-+      inputs, labels = inputs.to(xla.device()), labels.to(xla.device())
++    with torch_xla.step():
++      inputs, labels = inputs.to(torch_xla.device()), labels.to(torch_xla.device())
        optimizer.zero_grad()
        outputs = ddp_model(inputs)
        loss = loss_fn(outputs, labels)
@@ -213,7 +213,7 @@ If you're using `DistributedDataParallel`, make the following changes:
 
  if __name__ == '__main__':
 -  mp.spawn(_mp_fn, args=(), nprocs=world_size)
-+  xla.launch(_mp_fn, args=())
++  torch_xla.launch(_mp_fn, args=())
 ```
 
 Additional information on PyTorch/XLA, including a description of its semantics

docs/source/contribute/configure-environment.md

@@ -95,12 +95,12 @@ pip install numpy torch torch_xla[tpu] \
 Create a file `test.py`:
 
 ``` python
-import torch_xla as xla
+import torch_xla
 
 # Optional
-xla.runtime.set_device_type("TPU")
+torch_xla.runtime.set_device_type("TPU")
 
-print("XLA devices:", xla.real_devices())
+print("XLA devices:", torch_xla.real_devices())
 ```
 
 Run the test script from your terminal:

test/test_devices.py

@@ -4,7 +4,7 @@
 import torch
 from torch import nn
 from torch.utils.data import TensorDataset, DataLoader
-import torch_xla as xla
+import torch_xla
 import torch_xla.core.xla_model as xm
 import torch_xla.runtime as xr
 import torch_xla.debug.metrics as met
@@ -24,35 +24,35 @@ def setUp(self):
                             (0, torch.device('xla:0')),
                             (3, torch.device('xla:3')))
   def test_device(self, index, expected):
-    device = xla.device(index)
+    device = torch_xla.device(index)
     self.assertEqual(device, expected)
 
   def test_devices(self):
-    self.assertEqual(xla.devices(),
+    self.assertEqual(torch_xla.devices(),
                      [torch.device(f'xla:{i}') for i in range(4)])
 
   def test_real_devices(self):
-    self.assertEqual(xla.real_devices(), [f'CPU:{i}' for i in range(4)])
+    self.assertEqual(torch_xla.real_devices(), [f'CPU:{i}' for i in range(4)])
 
   def test_device_count(self):
-    self.assertEqual(xla.device_count(), 4)
+    self.assertEqual(torch_xla.device_count(), 4)
 
   def test_sync(self):
-    torch.ones((3, 3), device=xla.device())
-    xla.sync()
+    torch.ones((3, 3), device=torch_xla.device())
+    torch_xla.sync()
 
     self.assertEqual(met.counter_value('MarkStep'), 1)
 
   def test_step(self):
-    with xla.step():
-      torch.ones((3, 3), device=xla.device())
+    with torch_xla.step():
+      torch.ones((3, 3), device=torch_xla.device())
 
     self.assertEqual(met.counter_value('MarkStep'), 2)
 
   def test_step_exception(self):
     with self.assertRaisesRegex(RuntimeError, 'Expected error'):
-      with xla.step():
-        torch.ones((3, 3), device=xla.device())
+      with torch_xla.step():
+        torch.ones((3, 3), device=torch_xla.device())
         raise RuntimeError('Expected error')
 
     self.assertEqual(met.counter_value('MarkStep'), 2)
@@ -69,7 +69,7 @@ def __init__(self):
       def forward(self, x):
         return self.linear(x)
 
-    model = TrivialModel().to(xla.device())
+    model = TrivialModel().to(torch_xla.device())
 
     batch_size = 16
     num_samples = 100
@@ -85,8 +85,9 @@ def forward(self, x):
     optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
 
     for inputs, labels in loader:
-      with xla.step():
-        inputs, labels = inputs.to(xla.device()), labels.to(xla.device())
+      with torch_xla.step():
+        inputs, labels = inputs.to(torch_xla.device()), labels.to(
+            torch_xla.device())
         optimizer.zero_grad()
         outputs = model(inputs)
         loss = loss_fn(outputs, labels)

test/tpu/tpu_info/test_cli.py

@@ -4,7 +4,7 @@
 import multiprocessing
 import os
 from typing import Dict, Optional
-import torch_xla as xla
+import torch_xla
 import torch_xla.runtime as xr
 import torch_xla.distributed.xla_multiprocessing as xmp
 from tpu_info import cli, device, metrics
@@ -30,7 +30,7 @@ def _init_tpu_and_wait(
   ):
     if env:
       os.environ.update(**env)
-    xla.device()
+    torch_xla.device()
     q.put(os.getpid())
     done.wait()