nit: Apply suggestions from code review

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com>

Fix type on fused attn tests

Signed-off-by: Kshitij  Janardan Lakhani <klakhani@login-preos01.a51.clusters.nvidia.com>

Author: KshitijLakhani

tests/jax/test_fused_attn.py

@@ -776,7 +776,7 @@ def to_dp_shardings(x):
         self.seq_length_offset_pspec = PartitionSpec(self.mesh_resource.dp_resource, None)
         self.seq_length_offset_sharding = NamedSharding(self.mesh, self.seq_length_offset_pspec)
 
-    def _test_forward(self):
+    def test_forward(self):
         """
         Test forward with JITted primitive and unJITted reference
         """
@@ -1150,7 +1150,7 @@ class TestFusedAttn:
             pytest.param(AttnBiasType.POST_SCALE_BIAS, BiasShape._11SS, id="POST_SCALE_BIAS-11SS"),
         ],
     )
-    def test_forward(
+    def _test_forward(
         b,
         s_q,
         s_kv,

transformer_engine/jax/cpp_extensions/attention.py

@@ -1306,6 +1306,7 @@ def check_supported(self):
                 f" {','.join(map(str, allowed_masks))} got: {self.config.attn_mask_type}"
             )
         # TODO: For now do not all  CP + AG + THD + Striped with NO_MASK
+        # TODO: For now do not allow CP + AG + THD + Striped with NO_MASK
         if self.config.attn_mask_type is AttnMaskType.NO_MASK and self.config.qkv_layout.is_thd():
             raise ValueError(f"{header} only supports CAUSAL_MASK for THD types")
 
@@ -1380,7 +1381,7 @@ def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig:
         )
 
     def all_gather_kv(self, k, v):
-        """Performs aa all-gather of k and v over context parallel ranks."""
+        """Performs an all-gather of k and v over context parallel ranks."""
 
         def ag(x):
             x = lax_paral_op(
@@ -1402,7 +1403,7 @@ def ag(x):
         return k, v  # fall through
 
     def all_gather_segment_ids_and_pos(self, kv_segment_ids, kv_segment_pos):
-        """Performs aa all-gather of kv segment ids and kv segment pos over context parallel ranks."""
+        """Performs an all-gather of kv segment ids and kv segment pos over context parallel ranks."""
         kv_segment_ids = lax_paral_op(
             kv_segment_ids, lax.all_gather, self.config.cp_axis, mesh=self.mesh, axis=1, tiled=True
         )