keras-team/keras · error · ValueError

Expected `{t_name}` to have shape (B, 1, T, S) or (B, N, T,

Error message

Expected `{t_name}` to have shape (B, 1, T, S) or (B, N, T, S) with N={N} but got {t.shape}.

What it means

Error "Expected `{t_name}` to have shape (B, 1, T, S) or (B, N, T, S) with N={N} but got {t.shape}." thrown in keras-team/keras.

Source

Thrown at keras/src/backend/jax/nn.py:1856

    # _dot_product_attention_xla
    B, T, N, H = query.shape
    G = N // K
    query = jnp.reshape(query, (B, T, K, G, H))

    def _reshape_to_grouped(t, t_name):
        if t is not None:
            while t.ndim < 4:
                if t.ndim == 3 and t.shape[1] == N:
                    t = jnp.expand_dims(t, axis=2)
                else:
                    t = jnp.expand_dims(t, axis=1)
            tB, tN, tT, tS = t.shape
            if tN == 1:
                t = jnp.broadcast_to(t[:, :, None, :, :], (tB, tN, G, tT, tS))
            else:
                if tN != N:
                    raise ValueError(
                        f"Expected `{t_name}` to have shape (B, 1, T, S) or "
                        f"(B, N, T, S) with N={N} but got {t.shape}."
                    )
                t = jnp.reshape(t, (tB, K, G, tT, tS))
        return t

    bias = _reshape_to_grouped(bias, "bias")
    mask = _reshape_to_grouped(mask, "mask")
    vmapped_fn = jax.vmap(
        _dot_product_attention_core,
        in_axes=(3, None, None, 2, 2, None, None),
        out_axes=3,
    )
    encoded = vmapped_fn(query, key, value, bias, mask, is_causal, scale)
    return jnp.reshape(encoded, output_shape)


def unfold(input, kernel_size, dilation=1, padding=0, stride=1):

View on GitHub (pinned to 7a34a03db6)

When it happens

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