Update

[ghstack-poisoned]
pytorch · Jan 23, 2025 · 2f15cc1 · 2f15cc1
1 parent 0d8cdd9
commit 2f15cc1
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 8 deletions.
diff --git a/torchao/dtypes/__init__.py b/torchao/dtypes/__init__.py
@@ -18,8 +18,8 @@
     MarlinSparseLayout,
     SemiSparseLayout,
     TensorCoreTiledLayout,
-    to_marlinqqq_quantized_intx,
     UintxLayout,
+    to_marlinqqq_quantized_intx,
 )
 from .utils import Layout, PlainLayout
 

diff --git a/torchao/dtypes/affine_quantized_tensor_ops.py b/torchao/dtypes/affine_quantized_tensor_ops.py
@@ -35,11 +35,11 @@
     _linear_fp_act_int4_weight_sparse_marlin_impl,
 )
 from torchao.dtypes.uintx.plain_layout import (
+    PlainAQTTensorImpl,
     _linear_fp_act_int8_weight_check,
     _linear_fp_act_int8_weight_impl,
     _linear_int8_act_int8_weight_check,
     _linear_int8_act_int8_weight_impl,
-    PlainAQTTensorImpl,
 )
 from torchao.dtypes.uintx.semi_sparse_layout import (
     _linear_int8_act_int8_weight_semi_structured_sparse_check,

diff --git a/torchao/dtypes/floatx/__init__.py b/torchao/dtypes/floatx/__init__.py
@@ -1,10 +1,10 @@
 from .float8_layout import (
+    Float8Layout,
+    Float8QuantizedTensor,
     _linear_fp8_act_fp8_weight_check,
     _linear_fp8_act_fp8_weight_impl,
     _linear_fp_act_fp8_weight_check,
     _linear_fp_act_fp8_weight_impl,
-    Float8Layout,
-    Float8QuantizedTensor,
     to_affine_quantized_float8,
 )
 from .floatx_tensor_core_layout import (
@@ -13,7 +13,6 @@
     to_scaled_tc_floatx,
 )
 
-
 __all__ = [
     "FloatxTensorCoreLayout",
     "Float8Layout",

diff --git a/torchao/dtypes/floatx/float8_layout.py b/torchao/dtypes/floatx/float8_layout.py
@@ -11,17 +11,17 @@
     AffineQuantizedTensor,
     register_layout,
 )
-from torchao.dtypes.utils import AQTTensorImpl, get_out_shape, Layout
+from torchao.dtypes.utils import AQTTensorImpl, Layout, get_out_shape
 from torchao.float8.inference import (
+    Float8MMConfig,
     _is_rowwise_scaled,
     addmm_float8_unwrapped_inference,
-    Float8MMConfig,
     preprocess_data,
 )
 from torchao.quantization.quant_primitives import (
+    FP8_TYPES,
     choose_qparams_affine_float8,
     dequantize_affine_float8,
-    FP8_TYPES,
     quantize_affine_float8,
 )
 from torchao.utils import _is_float8_type, fill_defaults