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

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

tests/jax/test_distributed_softmax.py

@@ -138,7 +138,11 @@ def impl_test_softmax(
     @pytest.mark.parametrize("data_shape", [[32, 12, 128, 128], [8, 8, 1024, 1024]])
     @pytest.mark.parametrize(
         "softmax_fusion",
-        [SoftmaxFusion.SCALED, SoftmaxFusion.SCALED_MASKED, SoftmaxFusion.SCALED_UPPER_TRIANG_MASKED],
+        [
+            SoftmaxFusion.SCALED,
+            SoftmaxFusion.SCALED_MASKED,
+            SoftmaxFusion.SCALED_UPPER_TRIANG_MASKED,
+        ],
     )
     @pytest.mark.parametrize("scale_factor", [1.0, 3.0])
     @pytest.mark.parametrize("dtype", DTYPES)

tests/jax/utils.py

@@ -267,7 +267,7 @@ def __call__(
             # softmax_offset shape: [1, h, 1, 1], attn_weights shape: [b, h, q, k]
             extra_col = jnp.broadcast_to(
                 softmax_offset,
-                (attn_weights.shape[0], softmax_offset.shape[1], attn_weights.shape[2], 1)
+                (attn_weights.shape[0], softmax_offset.shape[1], attn_weights.shape[2], 1),
             )
             attn_weights = jnp.concatenate([attn_weights, extra_col], axis=-1)
 
@@ -1043,14 +1043,14 @@ def __call__(self, qlen, klen, bidirectional=True):
 
 def convert_softmax_type_str_to_enum(softmax_type: str) -> AttnSoftmaxType:
     """Convert softmax_type string to AttnSoftmaxType enum.
-    
+
     Args:
         softmax_type: String representation of softmax type.
             One of "vanilla", "off_by_one", "learnable".
-    
+
     Returns:
         AttnSoftmaxType enum value.
-    
+
     Raises:
         ValueError: If softmax_type is not a valid option.
     """
@@ -1063,7 +1063,7 @@ def convert_softmax_type_str_to_enum(softmax_type: str) -> AttnSoftmaxType:
     else:
         raise ValueError(
             f"Unknown softmax_type: {softmax_type}. "
-            f"Valid options: 'vanilla', 'off_by_one', 'learnable'"
+            "Valid options: 'vanilla', 'off_by_one', 'learnable'"
         )
 
 

transformer_engine/jax/flax/module.py

@@ -213,8 +213,8 @@ def __call__(self, inputs: Array, mask: Array = None, bias: Array = None) -> jnp
                 outputs = jax_scaled_upper_triang_masked_softmax(logits, self.scale_factor)
             else:
                 raise ValueError(
-                    f"Unsupported softmax fusion: {self.softmax_fusion}. softmax_fusion must be [SCALED,"
-                    " SCALED_MASKED, SCALED_UPPER_TRIANG_MASKED]"
+                    f"Unsupported softmax fusion: {self.softmax_fusion}. softmax_fusion must be"
+                    " [SCALED, SCALED_MASKED, SCALED_UPPER_TRIANG_MASKED]"
                 )
         assert input_dtype == outputs.dtype
         return outputs

transformer_engine/jax/flax/transformer.py

@@ -236,13 +236,15 @@ def apply_swa_mask(original_mask: Array) -> Array:
             # softmax_offset shape: [1, h, 1, 1], attn_weights shape: [b, h, q, k]
             extra_col = jnp.broadcast_to(
                 softmax_offset,
-                (attn_weights.shape[0], softmax_offset.shape[1], attn_weights.shape[2], 1)
+                (attn_weights.shape[0], softmax_offset.shape[1], attn_weights.shape[2], 1),
             )
             attn_weights = jnp.concatenate([attn_weights, extra_col], axis=-1)
 
             # Pad mask if present to match new shape
             if mask is not None:
-                mask = jnp.pad(mask, ((0, 0), (0, 0), (0, 0), (0, 1)), mode='constant', constant_values=0)
+                mask = jnp.pad(
+                    mask, ((0, 0), (0, 0), (0, 0), (0, 1)), mode="constant", constant_values=0
+                )
 
         def convert_to_softmax_type(attn_mask_type, mask):
             """Convert the attn_mask_type to SoftmaxFusion"""