def wrap_flash_attention(
    query,
    key,
    value,
    decoder_segment_ids,
    custom_mask=None,
    attn_logits_soft_cap=None,
    head_shards=1,
    q_seq_shards=1,
):
    """Applies a wrapped flash attention mechanism using the Splash kernel.
    This function prepares the appropriate attention mask (causal or custom),
    constructs a multi-head mask, and applies the Splash multi-head attention
    kernel to the provided query, key, and value tensors. It supports optional
    sharding and soft capping of attention logits.
    Args:
        query: jax.Array. The query tensor of shape
            (batch, num_heads, seq_len, head_dim).
        key: jax.Array. The key tensor of shape
            (batch, num_heads, seq_len, head_dim).
        value: jax.Array. The value tensor of shape
            (batch, num_heads, seq_len, head_dim).
        decoder_segment_ids: Optional. Segment IDs for the decoder, used for
            sharding or masking.
        custom_mask: Optional[jax.Array]. A custom attention mask to apply. If
            None, a causal mask is used.
        attn_logits_soft_cap: Optional[float]. If provided, applies a soft cap
            to the attention logits.
        head_shards: int, default=1. Number of shards for the attention heads.
        q_seq_shards: int, default=1. Number of shards for the query sequence
            dimension.
    Returns:
        jax.Array: The result of applying the Splash multi-head attention
            kernel to the inputs.
    Raises:
        AssertionError: If sharding along the sequence dimension is attempted
            with decoder_segment_ids.
    """
    if decoder_segment_ids is not None:
        assert query.shape[2] == decoder_segment_ids.q.shape[1], (
            "Sharding along sequence dimension not allowed"
            " in TPU kernel attention"
        )

    if custom_mask is not None:
        mask = splash_attention_mask.NumpyMask(array=custom_mask)
    else:
        mask = splash_attention_mask.CausalMask(
            shape=(query.shape[2], query.shape[2])
        )

    # Create multi-head mask
    multi_head_mask = splash_attention_mask.MultiHeadMask(
        masks=(mask,) * query.shape[1]
    )
    splash_kernel = splash_attention_kernel.make_splash_mha(
        mask=multi_head_mask,
        head_shards=head_shards,
        q_seq_shards=q_seq_shards,
        attn_logits_soft_cap=attn_logits_soft_cap,
    )

    return jax.vmap(splash_kernel)(
        query, key, value, segment_ids=decoder_segment_ids
    )
Applies a wrapped flash attention mechanism using the Splash kernel.
    This function prepares the appropriate attention mask (causal or custom),
    constructs a multi-head mask, and applies the Splash multi-head attention
    kernel to the provided query, key, and value tensors. It supports optional
    sharding and soft capping of attention logits.
    Args:
        query: jax.Array. The query tensor of shape
            (batch, num_heads, seq_len, head_dim).
        key: jax.Array. The key tensor of shape
            (batch, num_heads, seq_len, head_dim).
        value: jax.Array. The value tensor of shape
            (batch, num_heads, seq_len, head_dim).
        decoder_segment_ids: Optional. Segment IDs for the decoder, used for
            sharding or masking.
        custom_mask: Optional[jax.Array]. A custom attention mask to apply. If
            None, a causal mask is used.
        attn_logits_soft_cap: Optional[float]. If provided, applies a soft cap
            to the attention logits.
        head_shards: int, default=1. Number of shards for the attention heads.
        q_seq_shards: int, default=1. Number of shards for the query sequence
            dimension.
    Returns:
        jax.Array: The result of applying the Splash multi-head attention
            kernel to the inputs.
    Raises:
        AssertionError: If sharding along the sequence dimension is attempted
            with decoder_segment_ids.
    
keras/src/backend/jax/nn.py