nit: Fix the shape check for assert

KshitijLakhani · KshitijLakhani · commit 35d6d0f77db5 · 2026-02-23T17:27:12.000-08:00
Signed-off-by: Kshitij Lakhani &lt;klakhani@nvidia.com&gt;
diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
@@ -657,7 +657,7 @@ def batcher(batched_args, batch_dims, *, config):
                     "segment_ids must have more dims than segment_pos when adding batch dims; "
                     f"got segment_ids.ndim={segment_ids.ndim}, segment_pos.ndim={segment_pos.ndim}"
                 )
-                assert segment_ids.shape[segment_pos.ndim :] == segment_pos.shape, (
+                assert segment_ids.shape[-segment_pos.ndim :] == segment_pos.shape, (
                     "segment_pos must have same trailing shape as segment_ids when adding batch dims; "
                     f"got segment_ids.shape={segment_ids.shape}, segment_pos.shape={segment_pos.shape}"
                 )
@@ -1154,7 +1154,7 @@ def batcher(batched_args, batch_dims, *, config):
                     "segment_ids must have more dims than segment_pos when adding batch dims; "
                     f"got segment_ids.ndim={segment_ids.ndim}, segment_pos.ndim={segment_pos.ndim}"
                 )
-                assert segment_ids.shape[segment_pos.ndim :] == segment_pos.shape, (
+                assert segment_ids.shape[-segment_pos.ndim :] == segment_pos.shape, (
                     "segment_pos must have same trailing shape as segment_ids when adding batch dims; "
                     f"got segment_ids.shape={segment_ids.shape}, segment_pos.shape={segment_pos.shape}"
                 )

Original file line number	Diff line number	Diff line change
`@@ -657,7 +657,7 @@ def batcher(batched_args, batch_dims, *, config):`
`657`	`657`	`"segment_ids must have more dims than segment_pos when adding batch dims; "`
`658`	`658`	`f"got segment_ids.ndim={segment_ids.ndim}, segment_pos.ndim={segment_pos.ndim}"`
`659`	`659`	`)`
`660`		`- assert segment_ids.shape[segment_pos.ndim :] == segment_pos.shape, (`
	`660`	`+ assert segment_ids.shape[-segment_pos.ndim :] == segment_pos.shape, (`
`661`	`661`	`"segment_pos must have same trailing shape as segment_ids when adding batch dims; "`
`662`	`662`	`f"got segment_ids.shape={segment_ids.shape}, segment_pos.shape={segment_pos.shape}"`
`663`	`663`	`)`
`@@ -1154,7 +1154,7 @@ def batcher(batched_args, batch_dims, *, config):`
`1154`	`1154`	`"segment_ids must have more dims than segment_pos when adding batch dims; "`
`1155`	`1155`	`f"got segment_ids.ndim={segment_ids.ndim}, segment_pos.ndim={segment_pos.ndim}"`
`1156`	`1156`	`)`
`1157`		`- assert segment_ids.shape[segment_pos.ndim :] == segment_pos.shape, (`
	`1157`	`+ assert segment_ids.shape[-segment_pos.ndim :] == segment_pos.shape, (`
`1158`	`1158`	`"segment_pos must have same trailing shape as segment_ids when adding batch dims; "`
`1159`	`1159`	`f"got segment_ids.shape={segment_ids.shape}, segment_pos.shape={segment_pos.shape}"`
`1160`	`1160`	`)`