keras-team/keras · error · ValueError

`dot_product_attention` only supports 4D inputs. Received…

Error message

`dot_product_attention` only supports 4D inputs. Received: query.shape={query.shape}, key.shape={key.shape}, value.shape={value.shape}.

What it means

A rank guard in the NumPy `dot_product_attention`: after converting query, key and value to tensors, it verifies all three are 4D (batch, heads, sequence, feature dims), the layout the attention implementation assumes. If any of the three inputs has a rank other than 4 (e.g. a missing head or batch axis), this error reports all three shapes so you can see which tensor is malformed.

Solutions

1. Expand missing axes on query/key/value, e.g. query = query[None, ...] or keras.ops.expand_dims to add batch and/or head dimensions until each is (batch, heads, seq, features)
2. Split a 3D tensor into per-head 4D tensors if you folded the head axis into features
3. Pass tensors that already come from a MultiHeadAttention-style projection which produces 4D q/k/v

When it happens

Trigger: Thrown at keras/src/backend/numpy/nn.py:1435 when the library encounters an invalid state.

Common situations: See trigger scenarios.

AI-assisted analysis of keras-team/keras@7a34a03db6 (2026-08-25). Data as JSON: /api/errors/db888d808b9c6f78. Report an issue: GitHub.

Appendix: source

Thrown at keras/src/backend/numpy/nn.py:1435

    mask=None,
    scale=None,
    is_causal=False,
    flash_attention=None,
    attn_logits_soft_cap=None,
):
    if flash_attention is None:
        flash_attention = False
    if flash_attention:
        raise ValueError("Flash attention is not supported in numpy backend.")

    # Ref: jax.nn.dot_product_attention
    # https://github.com/jax-ml/jax/blob/jax-v0.4.32/jax/_src/nn/functions.py#L828
    # Not support `query_seq_lengths` and `key_value_seq_lengths` args
    query = convert_to_tensor(query)
    key = convert_to_tensor(key)
    value = convert_to_tensor(value)
    if len(query.shape) != 4:
        raise ValueError(
            "`dot_product_attention` only supports 4D inputs. "
            f"Received: query.shape={query.shape}, key.shape={key.shape}, "
            f"value.shape={value.shape}."
        )
    compute_dtype = backend.result_type(query.dtype, key.dtype, value.dtype)
    query = cast(query, compute_dtype)
    key = cast(key, compute_dtype)
    value = cast(value, compute_dtype)
    if bias is not None:
        bias = convert_to_tensor(bias, dtype=compute_dtype)

    _, _, _, H = key.shape
    scale = (1.0 / np.sqrt(H)) if scale is None else scale
    return _dot_product_attention_xla(
        query, key, value, bias, mask, is_causal, scale
    )

View on GitHub (pinned to 7a34a03db6)