[content] utils/nn.py (python)
class FlashMultiHeadAttention(nn.Module):
    """Custom MultiHeadAttention using jax.nn.dot_product_attention with flash attention support."""
    
    num_heads: int
    qkv_features: int
    dropout_rate: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool = True
    
    @nn.compact
    def __call__(self, inputs: jax.Array, mask: Optional[jax.Array] = None, deterministic: bool = False) -> jax.Array:
        """
        Args:
            inputs: Input array of shape (batch_size, seq_len, features)
            mask: Optional mask array of shape (batch_size, seq_len) or (seq_len, seq_len)
            deterministic: Whether to apply dropout
        """
        batch_size, seq_len, features = inputs.shape
        head_dim = self.qkv_features // self.num_heads
        
        # Linear projections for Q, K, V
        query = nn.Dense(
            self.qkv_features,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            name='query'
        )(inputs)
        key = nn.Dense(
            self.qkv_features,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            name='key'
        )(inputs)
        value = nn.Dense(
            self.qkv_features,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            name='value'
        )(inputs)
        
        # Reshape to (batch_size, seq_len, num_heads, head_dim)
        query = query.reshape(batch_size, seq_len, self.num_heads, head_dim)
        key = key.reshape(batch_size, seq_len, self.num_heads, head_dim)
        value = value.reshape(batch_size, seq_len, self.num_heads, head_dim)
        
        # Transpose to (batch_size, num_heads, seq_len, head_dim) for dot_product_attention
        query = query.transpose(0, 2, 1, 3)
        key = key.transpose(0, 2, 1, 3)
        value = value.transpose(0, 2, 1, 3)
        
        # Prepare mask for dot_product_attention
        attention_mask = None
        if mask is not None:
            # Convert mask to the format expected by dot_product_attention
            if mask.ndim == 2:
                # (seq_len, seq_len) -> (1, 1, seq_len, seq_len)
                attention_mask = mask[None, None, :, :]
            elif mask.ndim == 3:
                # (batch_size, seq_len, seq_len) -> (batch_size, 1, seq_len, seq_len)
                attention_mask = mask[:, None, :, :]
            else:
                # Assume it's already in the right format
                attention_mask = mask
        
        # Apply attention with flash attention if enabled
        implementation = 'cudnn' if self.use_flash_attention else 'xla'
        attention_output = jax.nn.dot_product_attention(
            query=query,
            key=key,
            value=value,
            mask=attention_mask,
            implementation=implementation
        )
        
        # Transpose back to (batch_size, seq_len, num_heads, head_dim)
        attention_output = attention_output.transpose(0, 2, 1, 3)
        
        # Reshape to (batch_size, seq_len, qkv_features)
        attention_output = attention_output.reshape(batch_size, seq_len, self.qkv_features)
        
        # Apply dropout if not deterministic
        if not deterministic and self.dropout_rate > 0.0:
            attention_output = nn.Dropout(
                rate=self.dropout_rate,
                deterministic=deterministic
            )(attention_output)
        
        # Final linear projection
        output = nn.Dense(
            features,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            name='output'
        )(attention_output)
        
        return output