Fix types for forward attention for JAX. (#704)

Bias and seed can both be None, type checking is failed otherwise.
Signed-off-by: George Karpenkov <george@metaworld.me>

transformer_engine/jax/cpp_extensions.py

@@ -2075,9 +2075,10 @@ class SelfFusedAttnFwdPrimitive(BasePrimitive):
 register_primitive(SelfFusedAttnFwdPrimitive)
 
 
-def self_fused_attn_fwd(qkv: jnp.ndarray, bias: jnp.ndarray, seqlen: jnp.ndarray, seed: jnp.ndarray,
-                        attn_bias_type: NVTE_Bias_Type, attn_mask_type: NVTE_Mask_Type,
-                        scaling_factor: float, dropout_probability: float, is_training: bool):
+def self_fused_attn_fwd(qkv: jnp.ndarray, bias: jnp.ndarray | None, seqlen: jnp.ndarray,
+                        seed: jnp.ndarray | None, attn_bias_type: NVTE_Bias_Type,
+                        attn_mask_type: NVTE_Mask_Type, scaling_factor: float,
+                        dropout_probability: float, is_training: bool):
     """
     Wrapper for TE self fused attention fwd
     Return BMM1 -> (PreScaleBias) -> Scale -> (PostScaleBias) -> Softmax -> (Dropout) -> BMM2

transformer_engine/jax/fused_attn.py

@@ -66,9 +66,10 @@ def is_fused_attn_kernel_available(q_type, kv_type, qkv_layout, attn_bias_type,
                            max_seqlen_q, max_seqlen_kv, head_dim).is_fused_attn_kernel_available()
 
 
-def self_fused_attn(qkv: jnp.ndarray, bias: jnp.ndarray, mask: jnp.ndarray, seed: jnp.ndarray,
-                    attn_bias_type: AttnBiasType, attn_mask_type: AttnMaskType,
-                    scaling_factor: float, dropout_probability: float, is_training: bool):
+def self_fused_attn(qkv: jnp.ndarray, bias: jnp.ndarray | None, mask: jnp.ndarray,
+                    seed: jnp.ndarray | None, attn_bias_type: AttnBiasType,
+                    attn_mask_type: AttnMaskType, scaling_factor: float,
+                    dropout_probability: float, is_training: bool):
     """
     Self fused attention wrapper
     """
@@ -86,19 +87,22 @@ def self_fused_attn(qkv: jnp.ndarray, bias: jnp.ndarray, mask: jnp.ndarray, seed
 
 
 @partial(jax.custom_vjp, nondiff_argnums=(4, 5, 6, 7, 8))
-def _self_fused_attn(qkv: jnp.ndarray, bias: jnp.ndarray, mask: jnp.ndarray, seed: jnp.ndarray,
-                     attn_bias_type: AttnBiasType, attn_mask_type: AttnMaskType,
-                     scaling_factor: float, dropout_probability: float, is_training: bool):
+def _self_fused_attn(qkv: jnp.ndarray, bias: jnp.ndarray | None, mask: jnp.ndarray,
+                     seed: jnp.ndarray | None, attn_bias_type: AttnBiasType,
+                     attn_mask_type: AttnMaskType, scaling_factor: float,
+                     dropout_probability: float, is_training: bool):
 
     output, _ = _self_fused_attn_fwd_rule(qkv, bias, mask, seed, attn_bias_type, attn_mask_type,
                                           scaling_factor, dropout_probability, is_training)
     return output
 
 
-def _self_fused_attn_fwd_rule(qkv: jnp.ndarray, bias: jnp.ndarray, mask: jnp.ndarray,
-                              seed: jnp.ndarray, attn_bias_type: AttnBiasType,
-                              attn_mask_type: AttnMaskType, scaling_factor: float,
-                              dropout_probability: float, is_training: bool):
+def _self_fused_attn_fwd_rule(qkv: jnp.ndarray, bias: jnp.ndarray | None,
+                              mask: jnp.ndarray, seed: jnp.ndarray | None,
+                              attn_bias_type: AttnBiasType,
+                              attn_mask_type: AttnMaskType,
+                              scaling_factor: float, dropout_probability: float,
+                              is_training: bool):
     if mask is None:
         batch, seqlen, *_ = qkv.shape
         actual_seqlen = jnp.full((batch,), seqlen, dtype=jnp.int32)