[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

transformer_engine/jax/attention.py

@@ -441,7 +441,11 @@ def _segment_ids_pos_to_seqlens_offsets(
     # This fast path avoids expanding the mask to Q * KV matrix and instead allows us to
     # examine only O(Q+KV) elements.
     # TODO(huah): This fast path does not work for CP + THD + unbalanced, need to fix later
-    if context_parallel_load_balanced and attn_mask_type.is_causal() and (window_size is None or window_size == (-1, -1)):
+    if (
+        context_parallel_load_balanced
+        and attn_mask_type.is_causal()
+        and (window_size is None or window_size == (-1, -1))
+    ):
         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
         )
@@ -529,7 +533,11 @@ def tree_unflatten(cls, aux_data, children):
         return cls(*children)
 
     def get_seqlens_and_offsets(
-        self, attn_mask_type, qkv_layout, window_size, max_segments_per_seq,
+        self,
+        attn_mask_type,
+        qkv_layout,
+        window_size,
+        max_segments_per_seq,
         context_parallel_load_balanced: bool = False,
     ):
         """
@@ -552,7 +560,7 @@ def get_seqlens_and_offsets(
                 attn_mask_type,
                 window_size,
                 max_segments_per_seq,
-                context_parallel_load_balanced
+                context_parallel_load_balanced,
             )
         else:
             q_seqlens, kv_seqlens = _segment_ids_to_seqlens(

transformer_engine/jax/cpp_extensions/attention.py

@@ -465,7 +465,7 @@ def impl(
                 config.qkv_layout,
                 config.window_size,
                 config.max_segments_per_seq,
-                config.context_parallel_load_balanced
+                config.context_parallel_load_balanced,
             )
         )
 
@@ -864,7 +864,7 @@ def impl(
                 config.qkv_layout,
                 config.window_size,
                 config.max_segments_per_seq,
-                config.context_parallel_load_balanced
+                config.context_parallel_load_balanced,
             )
         )