Apply ruff formatting fix

Signed-off-by: Benji Beck <benjibeck@meta.com>

Author: bbeckca

test/quantization/quantize_/workflows/float8/test_float8_semi_sparse.py

@@ -5,7 +5,9 @@
 # LICENSE file in the root directory of this source tree.
 
 import unittest
-from torchao.quantization.quantize_.workflows.float8.float8_semi_sparse_tensor import Float8SemiSparseTensor
+from torchao.quantization.quantize_.workflows.float8.float8_semi_sparse_tensor import (
+    Float8SemiSparseTensor,
+)
 from torchao.quantization.quantize_.workflows.float8.float8_tensor import Float8Tensor
 from torchao.float8.inference import Float8MMConfig
 import torch
@@ -44,20 +46,25 @@ def test_sparse_vs_dense_fp8(self, sizes):
         linear = torch.nn.Linear(K, N, dtype=dtype, device=device)
 
         apply_fake_sparsity(linear)
-        
+
         mm_config = Float8MMConfig(use_fast_accum=True)
-        input_fp8 = Float8Tensor.from_hp(input, float8_dtype=torch.float8_e4m3fn, mm_config=mm_config)
-        
-        weight_fp8 = Float8Tensor.from_hp(linear.weight.data, float8_dtype=torch.float8_e4m3fn, mm_config=mm_config)
+        input_fp8 = Float8Tensor.from_hp(
+            input, float8_dtype=torch.float8_e4m3fn, mm_config=mm_config
+        )
+
+        weight_fp8 = Float8Tensor.from_hp(
+            linear.weight.data, float8_dtype=torch.float8_e4m3fn, mm_config=mm_config
+        )
         dense_output = torch.nn.functional.linear(input_fp8, weight_fp8, linear.bias)
-        
+
         weight_sparse_fp8 = Float8SemiSparseTensor.from_hp(linear.weight.data, [1, K])
-        sparse_output = torch.nn.functional.linear(input_fp8, weight_sparse_fp8, linear.bias)
-        
-        torch.testing.assert_close(
-            dense_output, sparse_output, atol=3e-1, rtol=3e-1
+        sparse_output = torch.nn.functional.linear(
+            input_fp8, weight_sparse_fp8, linear.bias
         )
 
+        torch.testing.assert_close(dense_output, sparse_output, atol=3e-1, rtol=3e-1)
+
+
 instantiate_parametrized_tests(TestFloat8SemiSparseTensor)
 
 

torchao/quantization/quantize_/workflows/float8/float8_semi_sparse_tensor.py

@@ -86,7 +86,7 @@ def from_hp(
 @implements(aten.t.default)
 def _(func, types, args, kwargs):
     from torch.utils._python_dispatch import return_and_correct_aliasing
-    
+
     self = args[0]
     new = Float8SemiSparseTensor(
         sparse=self.sparse,
@@ -98,8 +98,10 @@ def _(func, types, args, kwargs):
 
 def _linear_fp8_semi_sparse(input_tensor, weight_tensor, bias):
     from torchao.ops import rowwise_scaled_linear_sparse_cutlass_f8f8
-    from torchao.quantization.quantize_.workflows.float8.float8_tensor import Float8Tensor
-    
+    from torchao.quantization.quantize_.workflows.float8.float8_tensor import (
+        Float8Tensor,
+    )
+
     if isinstance(input_tensor, Float8Tensor):
         input = input_tensor.qdata
         input_scale = input_tensor.scale
@@ -108,16 +110,16 @@ def _linear_fp8_semi_sparse(input_tensor, weight_tensor, bias):
         input = input_tensor.qdata
         input_scale = input_tensor.scale
         out_dtype = input_tensor.dtype
-    
+
     weight = weight_tensor.sparse
     weight_meta = weight_tensor.meta
     weight_scale = weight_tensor.scale
-    
+
     # Reshape input_scale if needed: kernel expects scale to match input shape minus last dim
     # For input [B, K], scale should be [B] not [B, 1]
     if input_scale.dim() > input.dim() - 1:
         input_scale = input_scale.squeeze(-1)
-    
+
     return rowwise_scaled_linear_sparse_cutlass_f8f8(
         input, input_scale, weight, weight_meta, weight_scale, bias, out_dtype
     )
@@ -130,7 +132,7 @@ def _(func, types, args, kwargs):
     else:  # aten.mm.default
         input_tensor, weight_tensor = args
         bias = None
-    
+
     return _linear_fp8_semi_sparse(input_tensor, weight_tensor, bias)
 
 

torchao/quantization/quantize_/workflows/float8/float8_tensor.py

@@ -256,7 +256,7 @@ def _(func, types, args, kwargs):
         args[1],
         args[2] if len(args) > 2 else None,
     )
-    
+
     # If weight is not Float8Tensor, return NotImplemented to allow weight's dispatch to handle it
     if not isinstance(weight_tensor, Float8Tensor):
         return NotImplemented