Apply code formatting with pre-commit

- Run black formatter on Python files
- Fix trailing whitespace
- Add missing newlines at end of files
- Format spacing around operators

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: lihongyang1990

transformer_engine/common/recipe/__init__.py

@@ -345,9 +345,18 @@ class Float8BlockScaling(Recipe):
     fp8_mha: bool = False
 
     def __post_init__(self) -> None:
-        assert self.x_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for x"
-        assert self.w_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for w"
-        assert self.grad_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for grad"
+        assert self.x_block_scaling_dim in [
+            1,
+            2,
+        ], "Only 1D or 2D blocks supported for x"
+        assert self.w_block_scaling_dim in [
+            1,
+            2,
+        ], "Only 1D or 2D blocks supported for w"
+        assert self.grad_block_scaling_dim in [
+            1,
+            2,
+        ], "Only 1D or 2D blocks supported for grad"
         assert not (
             self.x_block_scaling_dim == 2 and self.w_block_scaling_dim == 2
         ), "2D by 2D block gemm not supported."

transformer_engine/debug/features/api.py

@@ -135,7 +135,10 @@ def fp8_gemm_enabled(
 
         Union[bool, Tuple[bool, Optional[int]]] - default is (True, None)
         """
-        return True, None  # if it is false, fp8_gemm will be turned off. Otherwise nothing happens.
+        return (
+            True,
+            None,
+        )  # if it is false, fp8_gemm will be turned off. Otherwise nothing happens.
 
     def modify_tensor_enabled(
         self,
@@ -479,7 +482,11 @@ def call_feature(self, call, feat_config, layer_name, **kwargs):
         """
         if call.__name__ == "inspect_tensor":
             kwargs_copy = kwargs.copy()
-            for k in ["quantizer", "columnwise_quantized_tensor", "rowwise_quantized_tensor"]:
+            for k in [
+                "quantizer",
+                "columnwise_quantized_tensor",
+                "rowwise_quantized_tensor",
+            ]:
                 if k not in call.__code__.co_varnames:
                     kwargs_copy.pop(k)
         else:

transformer_engine/debug/features/log_fp8_tensor_stats.py

@@ -11,7 +11,9 @@
 import nvdlfw_inspect.api as debug_api
 
 
-from nvdlfw_inspect.debug_features.log_tensor_stats import LogTensorStats as BaseLogTensorStats
+from nvdlfw_inspect.debug_features.log_tensor_stats import (
+    LogTensorStats as BaseLogTensorStats,
+)
 from nvdlfw_inspect.registry import Registry, api_method
 
 from transformer_engine.debug.features.utils.stats_buffer import STATS_BUFFERS
@@ -21,11 +23,21 @@
     Float8CurrentScalingQuantizer,
 )
 from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer
-from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer
-from transformer_engine.debug.features.utils import get_reduction_params, next_enabled_iter
+from transformer_engine.pytorch.tensor.float8_blockwise_tensor import (
+    Float8BlockQuantizer,
+)
+from transformer_engine.debug.features.utils import (
+    get_reduction_params,
+    next_enabled_iter,
+)
 
 
-ALL_RECIPE_NAMES = ["fp8_delayed_scaling", "fp8_current_scaling", "mxfp8", "fp8_block_scaling"]
+ALL_RECIPE_NAMES = [
+    "fp8_delayed_scaling",
+    "fp8_current_scaling",
+    "mxfp8",
+    "fp8_block_scaling",
+]
 
 
 def _get_recipe_name(quantizer: Optional[Quantizer]):
@@ -47,7 +59,10 @@ def _get_new_quantizer(recipe_name, fp8_dtype):
         return Float8BlockQuantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True)
     if recipe_name == "fp8_current_scaling":
         return Float8CurrentScalingQuantizer(
-            fp8_dtype=fp8_dtype, device=torch.device("cuda"), rowwise=True, columnwise=True
+            fp8_dtype=fp8_dtype,
+            device=torch.device("cuda"),
+            rowwise=True,
+            columnwise=True,
         )
     if recipe_name == "mxfp8":
         return MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True)
@@ -214,9 +229,14 @@ def update_aux_dict(
         Needs to clean after usage, because it possibly change the usage of the quantized tensor.
         """
         fp8_dtype = None
-        if recipe_name in ["fp8_delayed_scaling", "fp8_current_scaling", "fp8_block_scaling"]:
+        if recipe_name in [
+            "fp8_delayed_scaling",
+            "fp8_current_scaling",
+            "fp8_block_scaling",
+        ]:
             assert isinstance(
-                quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer, Float8BlockQuantizer)
+                quantizer,
+                (Float8Quantizer, Float8CurrentScalingQuantizer, Float8BlockQuantizer),
             )
             fp8_dtype = quantizer.dtype
 
@@ -242,7 +262,8 @@ def update_aux_dict(
         finally:
             if isinstance(quantized_tensor, QuantizedTensor):
                 quantized_tensor.update_usage(
-                    rowwise_usage=old_rowwise_usage, columnwise_usage=old_columnwise_usage
+                    rowwise_usage=old_rowwise_usage,
+                    columnwise_usage=old_columnwise_usage,
                 )
 
     @api_method

transformer_engine/debug/features/log_tensor_stats.py

@@ -8,17 +8,26 @@
 
 import torch
 
-from nvdlfw_inspect.debug_features.log_tensor_stats import LogTensorStats as BaseLogTensorStats
+from nvdlfw_inspect.debug_features.log_tensor_stats import (
+    LogTensorStats as BaseLogTensorStats,
+)
 from nvdlfw_inspect.registry import Registry, api_method
 import nvdlfw_inspect.api as debug_api
 
 from transformer_engine.pytorch.tensor import QuantizedTensor, Quantizer
 from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor
 from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor
-from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage
-from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage
+from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import (
+    Float8TensorStorage,
+)
+from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import (
+    MXFP8TensorStorage,
+)
 from transformer_engine.debug.features.utils.stats_buffer import STATS_BUFFERS
-from transformer_engine.debug.features.utils import next_enabled_iter, get_reduction_params
+from transformer_engine.debug.features.utils import (
+    next_enabled_iter,
+    get_reduction_params,
+)
 
 
 @Registry.register_feature(namespace="transformer_engine")
@@ -90,7 +99,10 @@ class LogTensorStats(BaseLogTensorStats):
 
     def _get_supported_stats_list(self):
         """Returns stats this feature can log."""
-        return BaseLogTensorStats._get_supported_stats_list(None) | {"cur_amax", "dynamic_range"}
+        return BaseLogTensorStats._get_supported_stats_list(None) | {
+            "cur_amax",
+            "dynamic_range",
+        }
 
     @api_method
     def inspect_tensor_enabled(

transformer_engine/debug/features/utils/stats_buffer.py

@@ -131,7 +131,9 @@ def log(self):
             combiner = STATS[stat_name][1]
             stat_value = combiner(gathered_helper_stats)
             MetricLogger.log_scalar(
-                f"{self.layer_name}_{self.tensor_name}_{stat_name}", stat_value, self.iteration
+                f"{self.layer_name}_{self.tensor_name}_{stat_name}",
+                stat_value,
+                self.iteration,
             )
             output[(self.layer_name, self.tensor_name, stat_name, self.iteration)] = (
                 stat_value  # for debugging purposes
@@ -195,7 +197,13 @@ def reset(self):
         self.layers_to_next_iter: Dict[str, int] = {}
 
     def try_add_buffer(
-        self, layer_name, tensor_name, stats, options, reduction_group, reduce_within_microbatch
+        self,
+        layer_name,
+        tensor_name,
+        stats,
+        options,
+        reduction_group,
+        reduce_within_microbatch,
     ):
         """If buffer for such combination of stats/tensor_name/... is not present, this method creates it."""
         if (layer_name, tensor_name, options) in self.buffers:
@@ -205,7 +213,14 @@ def try_add_buffer(
         self.reduction_group_to_buffer[reduction_group].append(buffer)
 
     def feed(
-        self, layer_name, tensor_name, options, tensor, iteration, skip_reduction, aux_dict=None
+        self,
+        layer_name,
+        tensor_name,
+        options,
+        tensor,
+        iteration,
+        skip_reduction,
+        aux_dict=None,
     ):
         """
         Feeds the tensor into the respective buffer.

transformer_engine/debug/features/utils/stats_computation.py

@@ -122,15 +122,24 @@ def _get(buffers, stat_name):
 }
 
 STATS = {
-    "min": (lambda x, aux_dict: torch.min(x), lambda buffers: min(_get(buffers, "min"))),
-    "max": (lambda x, aux_dict: torch.max(x), lambda buffers: max(_get(buffers, "max"))),
-    "sum": (lambda x, aux_dict: torch.sum(x), lambda buffers: sum(_get(buffers, "sum"))),
+    "min": (
+        lambda x, aux_dict: torch.min(x),
+        lambda buffers: min(_get(buffers, "min")),
+    ),
+    "max": (
+        lambda x, aux_dict: torch.max(x),
+        lambda buffers: max(_get(buffers, "max")),
+    ),
+    "sum": (
+        lambda x, aux_dict: torch.sum(x),
+        lambda buffers: sum(_get(buffers, "sum")),
+    ),
     "mean": (
         lambda x, aux_dict: torch.mean(x),
         lambda buffers: sum(_get(buffers, "sum")) / sum(_get(buffers, "numel")),
     ),
     "numel": (
-        lambda x, aux_dict: x.numel() if hasattr(x, "numel") else x.get_data_tensors()[0].numel(),
+        lambda x, aux_dict: (x.numel() if hasattr(x, "numel") else x.get_data_tensors()[0].numel()),
         lambda buffers: sum(_get(buffers, "numel")),
     ),
     "l1_norm": (
@@ -151,7 +160,10 @@ def _get(buffers, stat_name):
             _get(buffers, "variance"), _get(buffers, "numel"), _get(buffers, "sum")
         ),
     ),
-    "cur_amax": (lambda x, aux_dict: x.abs().max(), lambda buffers: max(_get(buffers, "cur_amax"))),
+    "cur_amax": (
+        lambda x, aux_dict: x.abs().max(),
+        lambda buffers: max(_get(buffers, "cur_amax")),
+    ),
     "dynamic_range_top": (
         lambda x, aux_dict: _compute_dynamic_range_top(x),
         lambda buffers: max(_get(buffers, "dynamic_range_top")),

transformer_engine/debug/pytorch/debug_quantization.py

@@ -110,7 +110,9 @@ def get_plans_for_output_tensors(self) -> Tuple[bool, str]:
 
         inspect_tensor_enabled = self.process_enabled_api_call(
             debug_api.transformer_engine.inspect_tensor_enabled(
-                layer_name=self.layer_name, tensor_name=self.tensor_name, iteration=self.iteration
+                layer_name=self.layer_name,
+                tensor_name=self.tensor_name,
+                iteration=self.iteration,
             )
         )
 
@@ -135,7 +137,9 @@ def get_enabled_look_at_tensors(self):
 
         inspect_tensor_enabled = self.process_enabled_api_call(
             debug_api.transformer_engine.inspect_tensor_enabled(
-                layer_name=self.layer_name, tensor_name=self.tensor_name, iteration=self.iteration
+                layer_name=self.layer_name,
+                tensor_name=self.tensor_name,
+                iteration=self.iteration,
             )
         )
 
@@ -312,7 +316,10 @@ def quantize(
                 self.parent_quantizer.set_usage(rowwise=True)
 
         rowwise_gemm_tensor, columnwise_gemm_tensor = None, None
-        if STANDARD_FP8_QUANTIZE in [self.rowwise_tensor_plan, self.columnwise_tensor_plan]:
+        if STANDARD_FP8_QUANTIZE in [
+            self.rowwise_tensor_plan,
+            self.columnwise_tensor_plan,
+        ]:
             quantized_tensor = self.parent_quantizer(tensor)
             # if both rowwise_tensor_plan and columnwise_tensor_plan need to be in fp8,
             # one tensor with columnwise=True and rowwise=True is computed
@@ -542,7 +549,10 @@ def _update_parent_quantizer_usage(self):
             self.columnwise_usage and self.columnwise_tensor_plan == STANDARD_FP8_QUANTIZE
         )
 
-        if STANDARD_FP8_QUANTIZE in [self.rowwise_tensor_plan, self.columnwise_tensor_plan]:
+        if STANDARD_FP8_QUANTIZE in [
+            self.rowwise_tensor_plan,
+            self.columnwise_tensor_plan,
+        ]:
             self.parent_quantizer.set_usage(
                 rowwise=rowwise_gemm_quantize,
                 columnwise=columnwise_gemm_quantize,

transformer_engine/jax/attention.py

@@ -448,7 +448,10 @@ def run_length_fill_flattened(segment_ids_flattened) -> jnp.ndarray:
     # counts:                            [[2 3 1 1 1 5 3 0 0 0 0 0 0 0 0 0], [1 2 3 2 2 4 2 0 0 0 0 0 0 0 0 0]]
     # Returns segment_ids_run_length_1d: [[2 2 3 3 3 0 1 0 5 5 5 5 5 0 0 0], [1 0 0 3 3 3 0 0 2 2 4 4 4 4 0 0]]
     boundary = jnp.concatenate(
-        [jnp.broadcast_to(True, (1,)), segment_ids_flattened[1:] != segment_ids_flattened[:-1]]
+        [
+            jnp.broadcast_to(True, (1,)),
+            segment_ids_flattened[1:] != segment_ids_flattened[:-1],
+        ]
     )
     run_ids = jnp.cumsum(boundary) - 1
     # Each element could, in worst case, start a run
@@ -502,7 +505,11 @@ def _segment_ids_pos_to_seqlens_offsets(
         window_size == (-1, -1) and not attn_mask_type.is_bottom_right()
     ):
         return _segment_ids_pos_to_seqlens_offsets_fast_causal_path(
-            segment_ids_q, segment_ids_kv, segment_pos_q, segment_pos_kv, max_segments_per_seq
+            segment_ids_q,
+            segment_ids_kv,
+            segment_pos_q,
+            segment_pos_kv,
+            max_segments_per_seq,
         )
 
     # (1 = attend, 0 = masked)
@@ -622,7 +629,10 @@ def tree_flatten(self):
         """
         Flatten method to register as a pytree node
         """
-        return ((self.seqlens, self.seq_offsets, self.segment_ids, self.segment_pos), None)
+        return (
+            (self.seqlens, self.seq_offsets, self.segment_ids, self.segment_pos),
+            None,
+        )
 
     @classmethod
     def tree_unflatten(cls, aux_data, children):

transformer_engine/jax/checkpoint_policies.py

@@ -19,7 +19,10 @@
 
 def te_gemms_saveable(prim, *_, **__) -> bool:
     """Checkpoint policy for Transformer Engine GEMMs."""
-    is_te_gemm = prim in {GemmPrimitive.outer_primitive, GroupedGemmPrimitive.outer_primitive}
+    is_te_gemm = prim in {
+        GemmPrimitive.outer_primitive,
+        GroupedGemmPrimitive.outer_primitive,
+    }
     # Workaround to include JAX's scaled_matmul until JAX checkpoint policies for dots are
     # updated to include it.
     is_jax_scaled_matmul = prim.name == "scaled_matmul_wrapper"