formating config.py

Author: HDCharles

torchao/float8/config.py

@@ -99,13 +99,13 @@ def short_str(self):
 
     def __post_init__(self):
         if self.scaling_type is ScalingType.STATIC:
-            assert (
-                self.static_scale is not None
-            ), "static_scale must be specified for static scaling"
+            assert self.static_scale is not None, (
+                "static_scale must be specified for static scaling"
+            )
         if self.scaling_granularity is ScalingGranularity.AXISWISE:
-            assert (
-                self.scaling_type is ScalingType.DYNAMIC
-            ), "only dynamic scaling type is supported for axiswise scaling granularity"
+            assert self.scaling_type is ScalingType.DYNAMIC, (
+                "only dynamic scaling type is supported for axiswise scaling granularity"
+            )
         assert self.target_dtype is None or (
             self.target_dtype.is_floating_point and self.target_dtype.itemsize == 1
         ), "must specify a 8-bit floating-point dtype"
@@ -130,9 +130,9 @@ class DelayedScalingConfig:
     scale_fn_name: str = "max"
 
     def __post_init__(self):
-        assert (
-            self.scale_fn_name == "max"
-        ), f"{self.scale_fn_name} is not implemented yet. Only max is supported for now."
+        assert self.scale_fn_name == "max", (
+            f"{self.scale_fn_name} is not implemented yet. Only max is supported for now."
+        )
 
 
 @dataclass(frozen=True)
@@ -148,7 +148,6 @@ class Float8GemmConfig:
 
 # Pre-made recipes for common configurations
 class Float8LinearRecipeName(enum.Enum):
-
     # Default, dynamic per-tensor scaling with the cuBLAS tensorwise kernel
     TENSORWISE = "tensorwise"
 
@@ -291,7 +290,9 @@ def __post_init__(self):
 
         # float8 all-gather only supports tensorwise, in the future may support blockwise
         if self.cast_config_weight.scaling_granularity != ScalingGranularity.TENSORWISE:
-            assert not self.enable_fsdp_float8_all_gather, f"enable_fsdp_float8_all_gather only supports tensorwise scaling granularity, got {self.cast_config_weight.scaling_granularity}"
+            assert not self.enable_fsdp_float8_all_gather, (
+                f"enable_fsdp_float8_all_gather only supports tensorwise scaling granularity, got {self.cast_config_weight.scaling_granularity}"
+            )
 
         # save some characters in the compatibility checks below
         cc_i = self.cast_config_input
@@ -310,9 +311,9 @@ def __post_init__(self):
         ):
             is_disabled_1 = cc1.scaling_type is ScalingType.DISABLED
             is_disabled_2 = cc1.scaling_type is ScalingType.DISABLED
-            assert (
-                is_disabled_1 == is_disabled_2
-            ), f"incompatible operand precision for {gemm_name}"
+            assert is_disabled_1 == is_disabled_2, (
+                f"incompatible operand precision for {gemm_name}"
+            )
 
         for cc1, cc2, operand_name, default_dtype in [
             (cc_i, cc_i_gw, "input", e4m3_dtype),
@@ -324,9 +325,9 @@ def __post_init__(self):
                 object.__setattr__(cc1, "target_dtype", default_dtype)
             if cc2.target_dtype is None:
                 object.__setattr__(cc2, "target_dtype", default_dtype)
-            assert (
-                cc1.target_dtype == cc2.target_dtype
-            ), f"{operand_name} must be cast to the same dtype in both matmuls it's used in"
+            assert cc1.target_dtype == cc2.target_dtype, (
+                f"{operand_name} must be cast to the same dtype in both matmuls it's used in"
+            )
 
         # See the comments around `force_recompute_fp8_weight_in_bwd` for more details of this warning.
         if (
@@ -357,9 +358,9 @@ def from_recipe_name(
         """
         if type(recipe_name) == str:
             valid_names = [n.value for n in Float8LinearRecipeName]
-            assert (
-                recipe_name in valid_names
-            ), f"recipe_name {recipe_name} not in valid names {valid_names}"
+            assert recipe_name in valid_names, (
+                f"recipe_name {recipe_name} not in valid names {valid_names}"
+            )
             recipe_name = Float8LinearRecipeName(recipe_name)
 
         if recipe_name is Float8LinearRecipeName.TENSORWISE:
@@ -385,7 +386,6 @@ def from_recipe_name(
             )
 
         elif recipe_name is Float8LinearRecipeName.ROWWISE_WITH_GW_HP:
-
             # output_hp = input_fp8_axiswise_dim0 @ weight_t_axiswise_dim1
             cc_i = CastConfig(scaling_granularity=ScalingGranularity.AXISWISE)
             cc_w = CastConfig(scaling_granularity=ScalingGranularity.AXISWISE)