From fe9c61fd2300be66fc743a49696a19a813c76638 Mon Sep 17 00:00:00 2001 From: andrewwhitecdw Date: Wed, 12 Aug 2026 15:03:37 -0500 Subject: [PATCH] fix: comment typo 'mas' -> 'mask' in two TODOs Squashed to single commit for review. Original PR: https://github.com/andrewwhitecdw/TransformerEngine/pull/6 --- tests/jax/test_fused_attn.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 1dd92fe181..a657431fdc 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -957,7 +957,7 @@ def test_forward(self): customcall_args = [ # Put test data onto each GPU for distributed. - # TODO(mgoldfarb-nvidia): We will need to add reordering for bias, mas and + # TODO(mgoldfarb-nvidia): We will need to add reordering for bias, mask and # THD params once we support those features on CP. jax.device_put(self.cp_reorder_fn(self.q), self.qkvo_sharding), jax.device_put(self.cp_reorder_fn(self.k), self.qkvo_sharding), @@ -1089,7 +1089,7 @@ def grad_func( self.dropout_rng, ] customcall_args = [ - # TODO(mgoldfarb-nvidia): We will need to add reordering for bias, mas and + # TODO(mgoldfarb-nvidia): We will need to add reordering for bias, mask and # THD params once we support those features on CP. jax.device_put(self.cp_reorder_fn(self.q), self.qkvo_sharding), jax.device_put(self.cp_reorder_fn(self.k), self.qkvo_sharding),