[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

@@ -706,29 +706,30 @@ def get_seqlens_and_offsets(
         # something like : segment_ids (B, batch, seq), segment_pos (batch, seq)).
         if q_segment_ids.ndim < q_segment_pos.ndim or kv_segment_ids.ndim < kv_segment_pos.ndim:
             raise AssertionError(
-                "segment_ids must not have fewer dims than segment_pos; "
-                f"got q_segment_ids.ndim={q_segment_ids.ndim}, q_segment_pos.ndim={q_segment_pos.ndim}, "
-                f"kv_segment_ids.ndim={kv_segment_ids.ndim}, kv_segment_pos.ndim={kv_segment_pos.ndim}"
+                "segment_ids must not have fewer dims than segment_pos; got"
+                f" q_segment_ids.ndim={q_segment_ids.ndim},"
+                f" q_segment_pos.ndim={q_segment_pos.ndim},"
+                f" kv_segment_ids.ndim={kv_segment_ids.ndim},"
+                f" kv_segment_pos.ndim={kv_segment_pos.ndim}"
             )
         if not (
             q_segment_ids.shape[-q_segment_pos.ndim :] == q_segment_pos.shape
             and kv_segment_ids.shape[-kv_segment_pos.ndim :] == kv_segment_pos.shape
         ):
             raise AssertionError(
-                "segment_pos trailing shape must match segment_ids; "
-                f"got q_segment_ids.shape={q_segment_ids.shape}, q_segment_pos.shape={q_segment_pos.shape}, "
-                f"kv_segment_ids.shape={kv_segment_ids.shape}, kv_segment_pos.shape={kv_segment_pos.shape}"
+                "segment_pos trailing shape must match segment_ids; got"
+                f" q_segment_ids.shape={q_segment_ids.shape},"
+                f" q_segment_pos.shape={q_segment_pos.shape},"
+                f" kv_segment_ids.shape={kv_segment_ids.shape},"
+                f" kv_segment_pos.shape={kv_segment_pos.shape}"
             )
 
         if qkv_layout.is_thd():
             # THD: compute seqlens/offsets. Replicated segment_pos (more leading dims on segment_ids, e.g. if vmap)
             # i) Flatten leading batch dims so that segment_ids and segment_pos have the same number of leading dims,
-            # ii) vmap seqlens/offsets computation with segment_pos broadcast, 
+            # ii) vmap seqlens/offsets computation with segment_pos broadcast,
             # iii) reshape back to the original leading batch dims.
-            if (
-                q_segment_ids.ndim > q_segment_pos.ndim
-                or kv_segment_ids.ndim > kv_segment_pos.ndim
-            ):
+            if q_segment_ids.ndim > q_segment_pos.ndim or kv_segment_ids.ndim > kv_segment_pos.ndim:
                 n_batch_dims_q = q_segment_ids.ndim - q_segment_pos.ndim
                 n_batch_dims_kv = kv_segment_ids.ndim - kv_segment_pos.ndim
                 batch_shape_q = q_segment_ids.shape[:n_batch_dims_q]
@@ -738,9 +739,7 @@ def get_seqlens_and_offsets(
                 # assert flat_batch_q == flat_batch_kv, (
                 #     f"segment_ids batch size mismatch: {batch_shape_q} vs {batch_shape_kv}"
                 # )
-                q_flat = q_segment_ids.reshape(
-                    flat_batch_q, *q_segment_ids.shape[n_batch_dims_q:]
-                )
+                q_flat = q_segment_ids.reshape(flat_batch_q, *q_segment_ids.shape[n_batch_dims_q:])
                 kv_flat = kv_segment_ids.reshape(
                     flat_batch_kv, *kv_segment_ids.shape[n_batch_dims_kv:]
                 )
@@ -756,25 +755,23 @@ def single_batch(seg_id_q, seg_id_kv, seg_pos_q, seg_pos_kv):
                         max_segments_per_seq,
                     )
 
-                q_sl, kv_sl, q_off, kv_off = jax.vmap(
-                    single_batch, in_axes=(0, 0, None, None)
-                )(q_flat, kv_flat, q_segment_pos, kv_segment_pos)
+                q_sl, kv_sl, q_off, kv_off = jax.vmap(single_batch, in_axes=(0, 0, None, None))(
+                    q_flat, kv_flat, q_segment_pos, kv_segment_pos
+                )
 
                 q_seqlens = q_sl.reshape(*batch_shape_q, *q_sl.shape[1:])
                 kv_seqlens = kv_sl.reshape(*batch_shape_kv, *kv_sl.shape[1:])
                 q_offsets = q_off.reshape(*batch_shape_q, *q_off.shape[1:])
                 kv_offsets = kv_off.reshape(*batch_shape_kv, *kv_off.shape[1:])
             else:
-                q_seqlens, kv_seqlens, q_offsets, kv_offsets = (
-                    _segment_ids_pos_to_seqlens_offsets(
-                        q_segment_ids,
-                        kv_segment_ids,
-                        q_segment_pos,
-                        kv_segment_pos,
-                        attn_mask_type,
-                        window_size,
-                        max_segments_per_seq,
-                    )
+                q_seqlens, kv_seqlens, q_offsets, kv_offsets = _segment_ids_pos_to_seqlens_offsets(
+                    q_segment_ids,
+                    kv_segment_ids,
+                    q_segment_pos,
+                    kv_segment_pos,
+                    attn_mask_type,
+                    window_size,
+                    max_segments_per_seq,
                 )
         else:
             q_seqlens, kv_seqlens = _segment_ids_to_seqlens(

transformer_engine/jax/cpp_extensions/attention.py

@@ -638,7 +638,7 @@ def batcher(batched_args, batch_dims, *, config):
         assert FusedAttnFwdPrimitive.outer_primitive is not None
         q_bdim, _, _, _, _, seed_bdim, *_ = batch_dims
 
-        # When segment_ids are batched (vmap) and segment_pos are not, do not expand segment_pos to match. 
+        # When segment_ids are batched (vmap) and segment_pos are not, do not expand segment_pos to match.
         # The impl() layer treats segment_pos as replicated and computes seqlens/offsets per batch index
         # without materializing the full expanded segment_pos array.
         # Assert on invalid case (segment_ids.ndim < segment_pos.ndim)