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

pre-commit-ci[bot] · pre-commit-ci[bot] · commit 9fc5a745635b · 2025-11-26T07:22:27.000Z
for more information, see https://pre-commit.ci
diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
@@ -1676,7 +1676,7 @@ def kv_seqoffsets_for_striped_for_rank(
         max_segments_per_seq,
     ):
         """Extract the kv seqoffsets for striped primitive (post AG) from the sharded kv seg ids and seg pos,
-           AG kv seg ids and seg pos."""
+        AG kv seg ids and seg pos."""
         # Calculate the segment pos change mask
         segment_changes_first_true = jnp.concatenate(
             [

Original file line number	Diff line number	Diff line change
`@@ -1676,7 +1676,7 @@ def kv_seqoffsets_for_striped_for_rank(`
`1676`	`1676`	`max_segments_per_seq,`
`1677`	`1677`	`):`
`1678`	`1678`	`"""Extract the kv seqoffsets for striped primitive (post AG) from the sharded kv seg ids and seg pos,`
`1679`		`- AG kv seg ids and seg pos."""`
	`1679`	`+ AG kv seg ids and seg pos."""`
`1680`	`1680`	`# Calculate the segment pos change mask`
`1681`	`1681`	`segment_changes_first_true = jnp.concatenate(`
`1682`	`1682`	`[`