format python

Author: yaochengji

test/test_mp_all_gather.py

@@ -43,7 +43,8 @@ def _mp_fn(index):
 
     # Testing with a single replica group, channel_id and use_global_device_ids
     ordinal_tensor = torch.tensor([index], dtype=torch.float).to(device)
-    result = xm.all_gather(ordinal_tensor, dim=0, channel_id=0, use_global_device_ids=True)
+    result = xm.all_gather(
+        ordinal_tensor, dim=0, channel_id=0, use_global_device_ids=True)
     xm.mark_step()
 
     cpu_result = result.cpu()

test/test_mp_reduce_scatter.py

@@ -29,8 +29,14 @@ def _mp_fn(index):
     assert res.cpu().allclose(expected)
     xm.rendezvous('test_reduce_scatter')
 
-    res = xm.reduce_scatter(xm.REDUCE_SUM, xrand, scale, scatter_dim,
-                            world_size, channel_id=0, use_global_device_ids=True)
+    res = xm.reduce_scatter(
+        xm.REDUCE_SUM,
+        xrand,
+        scale,
+        scatter_dim,
+        world_size,
+        channel_id=0,
+        use_global_device_ids=True)
     expected_world = xm.all_reduce(xm.REDUCE_SUM, xrand, scale)
     xm.mark_step()
 

torch_xla/_internal/tpu.py

@@ -200,7 +200,7 @@ def get_tpu_type() -> str:
     env = get_tpu_env()
   except requests.HTTPError as e:
     raise EnvironmentError('Failed to get TPU metadata') from e
-  
+
   match = re.search(r"^([^-]*)-", env[xenv.ACCELERATOR_TYPE])
   if match:
     return match.group(1)

torch_xla/core/xla_model.py

@@ -534,8 +534,8 @@ def all_gather(value: torch.Tensor,
                groups: Optional[List[List[int]]] = None,
                output: Optional[torch.Tensor] = None,
                pin_layout: bool = True,
-               channel_id = None,
-               use_global_device_ids = None) -> torch.Tensor:
+               channel_id=None,
+               use_global_device_ids=None) -> torch.Tensor:
   """Performs an all-gather operation along a given dimension.
 
   Args:
@@ -588,7 +588,8 @@ def all_gather(value: torch.Tensor,
       return output
 
     result = torch_xla._XLAC._xla_all_gather(value, dim, shard_count, groups or
-                                             [], pin_layout, channel_id, use_global_device_ids)
+                                             [], pin_layout, channel_id,
+                                             use_global_device_ids)
     return result
 
   # Now the input should be a list of Tensors.
@@ -875,8 +876,8 @@ def reduce_scatter(reduce_type: str,
                    output: Optional[Union[torch.Tensor,
                                           List[torch.Tensor]]] = None,
                    pin_layout: bool = True,
-                   channel_id = None,
-                   use_global_device_ids = None) -> torch.Tensor:
+                   channel_id=None,
+                   use_global_device_ids=None) -> torch.Tensor:
   """Performs a XLA `ReduceScatter()` operation on the input tensor.
 
   See: https://www.tensorflow.org/xla/operation_semantics#reducescatter

torch_xla/distributed/xla_multiprocessing.py

@@ -123,13 +123,13 @@ def run(self, fn):
     """
     with self._lock:
       return fn()
-    
+
+
 ###############################################################################
 #
 # The following is modified from JAX: https://github.com/jax-ml/jax/blob/main/jax/_src/mesh_utils.py
 #
 ###############################################################################
-    
 
 _TPU_V5P = "v5p"
 _TPU_V6E = "v6e"
@@ -174,10 +174,7 @@ def _v6e_create_replica_groups() -> List | None:
   return None
 
 
-device_kind_handler_dict: dict[
-    str,
-    Callable[..., List | None],
-] = {
+device_kind_handler_dict: dict[str, Callable[..., List | None],] = {
     _TPU_V5P: _v5p_create_replica_groups,
     _TPU_V6E: _v6e_create_replica_groups
 }