[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

transformer_engine/pytorch/cpu_offload.py

@@ -340,9 +340,8 @@ def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any:
 
             self.tensor_tag_to_state[tensor_tag] = tensor
 
-            if (
-                self.current_group < self.num_offload_group
-                and self.tensor_need_offloading_checker(tensor)
+            if self.current_group < self.num_offload_group and self.tensor_need_offloading_checker(
+                tensor
             ):
                 self.tensor_tag_to_buf[tensor_tag] = tensor
         else:

transformer_engine/pytorch/tensor/float8_blockwise_tensor.py

@@ -203,7 +203,10 @@ def make_empty(
         columnwise_scale_inv = None
         if self.columnwise_usage:
             columnwise_data = torch.empty(
-                self.get_columnwise_shape(shape), dtype=torch.uint8, device=device, pin_memory=pin_memory
+                self.get_columnwise_shape(shape),
+                dtype=torch.uint8,
+                device=device,
+                pin_memory=pin_memory,
             )
             columnwise_scale_shape = self.get_scale_shape(shape, columnwise=True)
             columnwise_scale_inv = torch.empty(

transformer_engine/pytorch/tensor/float8_tensor.py

@@ -97,7 +97,7 @@ def make_empty(
         device: Optional[torch.device] = None,
         requires_grad: bool = False,
         pin_memory: bool = False,
-        rowwise: bool = None,   
+        rowwise: bool = None,
         columnwise: bool = None,
     ) -> Float8Tensor:
         rowwise = rowwise if rowwise is not None else self.rowwise_usage
@@ -624,6 +624,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None):
             dst, src = args[0], args[1]
             # Just copy FP8 attrs if copying between Float8Tensors
             if isinstance(src, Float8Tensor) and isinstance(dst, Float8Tensor):
+
                 def copy_tensor(src, dst, tensor_name):
                     src_is_none = getattr(src, tensor_name) is None
                     dst_is_none = getattr(dst, tensor_name) is None
@@ -633,13 +634,14 @@ def copy_tensor(src, dst, tensor_name):
                         setattr(dst, tensor_name, None)
                     elif not src_is_none and not dst_is_none:
                         getattr(src, tensor_name).copy_(getattr(dst, tensor_name))
+
                 copy_tensor(src, dst, "_data")
                 copy_tensor(src, dst, "_scale_inv")
                 if src._transpose is not None and dst._data is not None:
                     dst._create_transpose()
                 copy_tensor(src, dst, "_transpose")
                 return
-                
+
         elif func in _ops_to_preserve_subclass_in_fsdp2:
             # Ops in the _ops_to_preserve_subclass_in_fsdp2 are recommened to return the same class instance to work fine with the torch fsdp2
             warnings.warn(

transformer_engine/pytorch/tensor/quantized_tensor.py

@@ -366,19 +366,30 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None):
         # View op
         if func == torch.ops.aten.view.default:
             raise NotImplementedError("{cls.__name__} class does not support tensor views")
-        
+
         # Empty like op
         if func == torch.ops.aten.empty_like.default:
             tensor = args[0]
-            quantizer = tensor._quantizer # TODO - pgadzinski look if this makes sense
+            quantizer = tensor._quantizer  # TODO - pgadzinski look if this makes sense
             dtype = kwargs.get("dtype", tensor.dtype)
             device = kwargs.get("device", tensor.device)
             shape = kwargs.get("shape", tensor.shape)
             pin_memory = kwargs.get("pin_memory", False)
-            rowwise = (getattr(tensor, "_data", None) is not None) or (getattr(tensor, "_rowwise_data", None) is not None)
-            columnwise = (getattr(tensor, "_transpose", None) is not None) or (getattr(tensor, "_columnwise_data", None) is not None)
-
-            return quantizer.make_empty(shape, dtype=dtype, device=device, pin_memory=pin_memory, rowwise=rowwise, columnwise=columnwise)
+            rowwise = (getattr(tensor, "_data", None) is not None) or (
+                getattr(tensor, "_rowwise_data", None) is not None
+            )
+            columnwise = (getattr(tensor, "_transpose", None) is not None) or (
+                getattr(tensor, "_columnwise_data", None) is not None
+            )
+
+            return quantizer.make_empty(
+                shape,
+                dtype=dtype,
+                device=device,
+                pin_memory=pin_memory,
+                rowwise=rowwise,
+                columnwise=columnwise,
+            )
 
         def maybe_unwrap(arg):
             if isinstance(arg, QuantizedTensor):