keras-team/keras · error · ValueError

Sharding along sequence dimension not allowed in TPU kernel

Error message

Sharding along sequence dimension not allowed in TPU kernel attention

What it means

Error "Sharding along sequence dimension not allowed in TPU kernel attention" thrown in keras-team/keras.

Source

Thrown at keras/src/backend/jax/ops/nn.py:1613

        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:
        if query.shape[2] != decoder_segment_ids.q.shape[1]:
            raise ValueError(
                "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,

View on GitHub (pinned to 7a34a03db6)

When it happens

Trigger: Thrown at keras/src/backend/jax/ops/nn.py:1613 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/5a8f9b07e62532d4. Report an issue: GitHub.