tf.keras.ops.dot_product_attention

Scaled dot product attention function.

tf.keras.ops.dot_product_attention(
    query,
    key,
    value,
    bias=None,
    mask=None,
    scale=None,
    is_causal=False,
    flash_attention=None,
    attn_logits_soft_cap=None
)

Computes the attention function on Q (query), K (key), and V(value): attention(Q, K, V) = softmax(Q * K / sqrt(d)) * V. If we define logits as the output of Q * K and the probs as the output of softmax.

Throughout this function, we utilize the following notation to represent the shape of array:

  • B: batch size
  • S: length of the key/value
  • T: length of the query
  • N: number of attention heads
  • H: dimensions of each attention head
  • K: number of key/value heads
  • G: number of groups, which equals to N // K

query The query array with the shape of (B, T, N, H).
key The key array with the shape of (B, S, K, H). When K equals N, multi-headed attention (MHA) is performed. Otherwise, grouped query attention (GQA) is performed if N is a multiple of K. and multi-query attention (MQA) is performed if K==1 (a special case of GQA).
value The value array with the same shape of key.
bias Optional bias array to be added to logits. The shape must be broadcastable to (B, N, T, S).
mask Optional mask array used to filter out logits. It is a boolean mask where True indicates the element should take part in attention. For an additive mask, users should pass it to bias. The shape must be broadcastable to (B, N, T, S).
scale Optional scale for the logits. If None, the scale will be set to 1.0 / sqrt(H).
is_causal Whether to apply causal mask.
flash_attention Whether to use flash attention. If None, it will attempt to use flash attention if the required conditions are met. Typically, the inputs must be in float16 and bfloat16 dtype and the input layout requirements may vary depending on the backend.
attn_logits_soft_cap The value limit for maximum value of the attention logits before the softmax function is applied. This is only supported in JAX TPU backend. Defaults to None.

An array of the attention output with the same shape of query.

Example:

query = keras.random.normal((2, 4, 8, 16))
key = keras.random.normal((2, 6, 8, 16))
value = keras.random.normal((2, 6, 8, 16))
keras.ops.nn.dot_product_attention(query, key, value).shape
(2, 4, 8, 16)