fix format issue

Author: zxd1997066

test/prototype/mx_formats/test_kernels.py

@@ -46,10 +46,10 @@
 from torchao.prototype.mx_formats.mx_tensor import ScaleCalculationMode, to_dtype, to_mx
 from torchao.prototype.mx_formats.utils import to_blocked
 from torchao.utils import (
+    get_current_accelerator_device,
     is_sm_at_least_89,
     is_sm_at_least_100,
     torch_version_at_least,
-    get_current_accelerator_device,
 )
 
 torch.manual_seed(0)

test/prototype/mx_formats/test_mx_tensor.py

@@ -27,10 +27,10 @@
 from torchao.prototype.mx_formats.utils import from_blocked, to_blocked
 from torchao.quantization.utils import compute_error
 from torchao.utils import (
+    get_current_accelerator_device,
     is_sm_at_least_89,
     is_sm_at_least_90,
     torch_version_at_least,
-    get_current_accelerator_device,
 )
 
 torch.manual_seed(2)
@@ -663,7 +663,10 @@ def test_to_blocked_from_blocked_roundtrip(shape, use_triton_kernel: bool):
 
 
 @pytest.mark.skipif(not torch.accelerator.is_available(), reason="GPU not available")
-@pytest.mark.skipif(torch.cuda.is_available() and not torch_version_at_least("2.8.0"), reason="requires PyTorch 2.8+")
+@pytest.mark.skipif(
+    torch.cuda.is_available() and not torch_version_at_least("2.8.0"),
+    reason="requires PyTorch 2.8+",
+)
 @pytest.mark.parametrize("transpose", [False, True])
 @pytest.mark.parametrize(
     "shape",
@@ -717,7 +720,10 @@ def test_scale_shape_matches_qdata(transpose, shape):
 
 
 @pytest.mark.skipif(not torch.accelerator.is_available(), reason="GPU not available")
-@pytest.mark.skipif(torch.cuda.is_available() and not torch_version_at_least("2.8.0"), reason="requires PyTorch 2.8+")
+@pytest.mark.skipif(
+    torch.cuda.is_available() and not torch_version_at_least("2.8.0"),
+    reason="requires PyTorch 2.8+",
+)
 @pytest.mark.parametrize("elem_dtype", (torch.float8_e4m3fn, torch.float4_e2m1fn_x2))
 @pytest.mark.parametrize("transpose", [False, True])
 @pytest.mark.parametrize(

torchao/prototype/mx_formats/kernels.py

@@ -38,6 +38,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 def get_bits(x: torch.Tensor) -> str:
     bits_per_byte = 8
     # Numpy has a nice function to get the string representation of binary.
@@ -591,7 +592,6 @@ def triton_f6_e3m2_to_bf16(x: torch.Tensor) -> torch.Tensor:
     assert x.is_cuda or x.device.type == "xpu"
     assert output.is_cuda or output.device.type == "xpu"
 
-
     n_mx_blocks = x.shape[0]
     grid = lambda meta: (triton.cdiv(n_mx_blocks, meta["BLOCK_SIZE_IN"]),)
     triton_f6_to_bf16_kernel[grid](