From 6aeb71e343c6ec49248132c5fcb031cbf5404207 Mon Sep 17 00:00:00 2001 From: andrewwhitecdw Date: Tue, 18 Aug 2026 16:40:56 -0500 Subject: [PATCH] fix: comment typo 'mas' -> 'mask' in two TODOs Signed-off-by: andrewwhitecdw --- 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 b9d4ad670e..352ab64a0d 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),