Replace device=torch_xla.device() with device="xla"

Author: ghpvnist

API_GUIDE.md

@@ -15,7 +15,7 @@ import torch
 import torch_xla
 import torch_xla.core.xla_model as xm
 
-t = torch.randn(2, 2, device=torch_xla.device())
+t = torch.randn(2, 2, device="xla")
 print(t.device)
 print(t)
 ```
@@ -32,8 +32,8 @@ PyTorch operations can be performed on XLA tensors just like CPU or CUDA tensors
 For example, XLA tensors can be added together:
 
 ```python
-t0 = torch.randn(2, 2, device=torch_xla.device())
-t1 = torch.randn(2, 2, device=torch_xla.device())
+t0 = torch.randn(2, 2, device="xla")
+t1 = torch.randn(2, 2, device="xla")
 print(t0 + t1)
 ```
 
@@ -46,7 +46,7 @@ print(t0.mm(t1))
 Or used with neural network modules:
 
 ```python
-l_in = torch.randn(10, device=torch_xla.device())
+l_in = torch.randn(10, device="xla")
 linear = torch.nn.Linear(10, 20).to(torch_xla.device())
 l_out = linear(l_in)
 print(l_out)
@@ -56,7 +56,7 @@ Like other device types, XLA tensors only work with other XLA tensors on the
 same device. So code like
 
 ```python
-l_in = torch.randn(10, device=torch_xla.device())
+l_in = torch.randn(10, device="xla")
 linear = torch.nn.Linear(10, 20)
 l_out = linear(l_in)
 print(l_out)

benchmarks/matmul_bench.py

@@ -39,10 +39,7 @@ def main():
   """
 
   xla_bench_fn = lambda fn: do_bench(
-      fn,
-      return_mode='min',
-      sync_fn=lambda: xm.wait_device_ops(),
-      device=torch_xla.device())
+      fn, return_mode='min', sync_fn=lambda: xm.wait_device_ops(), device="xla")
   ind_bench_fn = lambda fn: do_bench(
       fn,
       return_mode='min',
@@ -53,7 +50,7 @@ def main():
   for dtype in dtypes:
     for inductor_matmul, xla_matmul in zip(
         get_matmuls(device='cuda', dtype=dtype, backend='inductor'),
-        get_matmuls(device=torch_xla.device(), dtype=dtype, backend='openxla')):
+        get_matmuls(device="xla", dtype=dtype, backend='openxla')):
       ind_lhs_shape, ind_rhs_shape, ind_fn = inductor_matmul
       xla_lhs_shape, xla_rhs_shape, xla_fn = xla_matmul
       assert ind_lhs_shape == xla_lhs_shape, f"Expect matmul shapes to match for benchmarking. Mismatch lhs: {ind_lhs_shape}, rhs: {xla_rhs_shape}"

contrib/kaggle/distributed-pytorch-xla-basics-with-pjrt.ipynb

@@ -273,7 +273,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2024-01-10T19:30:33.219878Z",
@@ -318,12 +318,12 @@
    ],
    "source": [
     "def add_ones(i, lock):\n",
-    "    x = torch.ones((3, 3), device=torch_xla.device())\n",
+    "    x = torch.ones((3, 3), device=\"xla\")\n",
     "    y = x + x\n",
-    "    \n",
+    "\n",
     "    # Run graph to compute `y` before printing\n",
     "    torch_xla.sync()\n",
-    "    \n",
+    "\n",
     "    with lock:\n",
     "        print(i, y)\n",
     "\n",
@@ -340,7 +340,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2024-01-10T19:30:35.656796Z",
@@ -378,10 +378,10 @@
    "source": [
     "def gather_ids(i, lock):\n",
     "    # Create a tensor on each device with the device ID\n",
-    "    t = torch.tensor([i], device=torch_xla.device())\n",
+    "    t = torch.tensor([i], device=\"xla\")\n",
     "    with lock:\n",
     "        print(i, t)\n",
-    "    \n",
+    "\n",
     "    # Collect and concatenate the IDs\n",
     "    ts = xm.all_gather(t)\n",
     "    torch_xla.sync()\n",
@@ -402,7 +402,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2024-01-10T19:30:38.315927Z",
@@ -479,7 +479,7 @@
     "        loss.backward()\n",
     "\n",
     "        optimizer.step()\n",
-    "        \n",
+    "\n",
     "        # Run the pending graph\n",
     "        torch_xla.sync()\n",
     "\n",

docs/source/learn/_pjrt.md

@@ -377,7 +377,7 @@ def _all_gather(index: int):
   # No need to pass in `rank` or `world_size`
   dist.init_process_group('xla', init_method='xla://')
 
-  t = torch.tensor([index], dtype=torch.int32, device=torch_xla.device())
+  t = torch.tensor([index], dtype=torch.int32, device="xla")
   output = [torch.zeros_like(t) for _ in range(dist.get_world_size())]
   dist.all_gather(output, t)
 

docs/source/learn/pytorch-on-xla-devices.md

@@ -14,7 +14,7 @@ import torch
 import torch_xla
 import torch_xla.core.xla_model as xm
 
-t = torch.randn(2, 2, device=torch_xla.device())
+t = torch.randn(2, 2, device="xla")
 print(t.device)
 print(t)
 ```
@@ -32,8 +32,8 @@ tensors.
 For example, XLA tensors can be added together:
 
 ``` python
-t0 = torch.randn(2, 2, device=torch_xla.device())
-t1 = torch.randn(2, 2, device=torch_xla.device())
+t0 = torch.randn(2, 2, device="xla")
+t1 = torch.randn(2, 2, device="xla")
 print(t0 + t1)
 ```
 
@@ -46,7 +46,7 @@ print(t0.mm(t1))
 Or used with neural network modules:
 
 ``` python
-l_in = torch.randn(10, device=torch_xla.device())
+l_in = torch.randn(10, device="xla")
 linear = torch.nn.Linear(10, 20).to(torch_xla.device())
 l_out = linear(l_in)
 print(l_out)
@@ -56,7 +56,7 @@ Like other device types, XLA tensors only work with other XLA tensors on
 the same device. So code like
 
 ``` python
-l_in = torch.randn(10, device=torch_xla.device())
+l_in = torch.randn(10, device="xla")
 linear = torch.nn.Linear(10, 20)
 l_out = linear(l_in)
 print(l_out)

docs/source/learn/troubleshoot.md

@@ -32,8 +32,8 @@ vm:~$ export PJRT_DEVICE=TPU
 vm:~$ python3
 >>> import torch
 >>> import torch_xla.core.xla_model as xm
->>> t1 = torch.tensor(100, device=torch_xla.device())
->>> t2 = torch.tensor(200, device=torch_xla.device())
+>>> t1 = torch.tensor(100, device="xla")
+>>> t2 = torch.tensor(200, device="xla")
 >>> print(t1 + t2)
 tensor(300, device='xla:0')
 ```

examples/scan/scan_examples.py

@@ -18,8 +18,8 @@ def cumsum(accumulated, element):
     return accumulated, accumulated
 
   # 2) Define an initial carry and the input tensor.
-  init_sum = torch.tensor([0.0], device=torch_xla.device())
-  xs = torch.tensor([1.0, 2.0, 3.0], device=torch_xla.device())
+  init_sum = torch.tensor([0.0], device="xla")
+  xs = torch.tensor([1.0, 2.0, 3.0], device="xla")
   torch_xla.sync()
 
   # 3) Call `scan` with our combine function, initial carry, and input tensor.
@@ -40,15 +40,15 @@ def scan_example_pytree():
   #    - 'sum' to accumulate the sum of all seen values
   #    - 'count' to count how many values have been seen
   carry = {
-      'sum': torch.tensor([0.0], device=torch_xla.device()),
-      'count': torch.tensor([0.0], device=torch_xla.device())
+      'sum': torch.tensor([0.0], device="xla"),
+      'count': torch.tensor([0.0], device="xla")
   }
 
   # 2) Define our input PyTree, which in this case is just a dictionary with one leaf:
   #    - 'values' is a 1D tensor representing data points we want to scan over.
   xs = {
       'values':
-          torch.arange(1, 6, dtype=torch.float32, device=torch_xla.device())
+          torch.arange(1, 6, dtype=torch.float32, device="xla")
   }
 
   # Here, xs['values'] has shape [5]. The `scan` function will automatically slice

test/ds/test_dynamic_shapes.py

@@ -163,7 +163,7 @@ def test_t_copy(self):
     self.assertEqual(t2_t.shape[1], 7)
 
   def test_nonzero_shape(self):
-    x = torch.tensor((0, 1, 2, 0, 3, 4), device=torch_xla.device())
+    x = torch.tensor((0, 1, 2, 0, 3, 4), device="xla")
     x_dim0_shape = torch_xla._XLAC._get_xla_tensor_dimension_size(
         torch.nonzero(x, as_tuple=False), 0)
     self.assertEqual(x_dim0_shape.item(), 4)
@@ -176,14 +176,14 @@ def test_nonzero_correctness(self):
     self.assertEqual(t2.cpu(), t2_aten)
 
   def test_masked_select_shape(self):
-    x = torch.tensor((0, 1, 2, 0, 3, 4), device=torch_xla.device())
+    x = torch.tensor((0, 1, 2, 0, 3, 4), device="xla")
     mask = x.ge(2)
     x_dim0_shape = torch_xla._XLAC._get_xla_tensor_dimension_size(
         torch.masked_select(x, mask), 0)
     self.assertEqual(x_dim0_shape.item(), 3)
 
   def test_nonzero_cast(self):
-    t1 = torch.ones(5, 2, device=torch_xla.device())
+    t1 = torch.ones(5, 2, device="xla")
     # Result of the nonzero should be the index type. Currently
     # index type is s64 on cpu and gpu, but s32 on TPU. We should be
     # able to cast it to any other type without error.

test/dynamo/test_dynamo.py

@@ -49,7 +49,7 @@ def inplace_update(self, a):
   def test_inplace_update_correctness(self, backend):
     dynamo_inplace = torch.compile(
         self.inplace_update, backend=backend, fullgraph=True)
-    t = torch.tensor([0, 1, 2], device=torch_xla.device())
+    t = torch.tensor([0, 1, 2], device="xla")
     for i in range(10):
       t = dynamo_inplace(t)
     self.assertTrue(torch.all(torch.eq(t.cpu(), torch.tensor([10, 11, 12]))))
@@ -131,7 +131,7 @@ def dummy_fn(self, a):
   def test_dynamo_with_trace(self):
     dynamo_dummy = torch.compile(
         self.dummy_fn, backend="openxla", fullgraph=True)
-    t = torch.randn(2, 3, 4, device=torch_xla.device())
+    t = torch.randn(2, 3, 4, device="xla")
     for i in range(10):
       with xp.Trace('build_graph'):
         t = dynamo_dummy(t)

test/metrics_compare_utils_test.py

@@ -275,7 +275,7 @@ def test_compare_metrics_reports_new_counters(self):
   def test_parse_real_metrics(self):
     print(
         'Testing against TPU. If this hangs, check that $XRT_TPU_CONFIG is set')
-    x = torch.rand(3, 5, device=torch_xla.device())
+    x = torch.rand(3, 5, device="xla")
     x = torch.flatten(x, 1)
     x = torch.roll(x, 1, 0)
     x = torch.flip(x, [0, 1])