jax-ml/jax · error · ValueError

{name} shape should be {shape}: but got {t.shape}

Error message

{name} shape should be {shape}: but got {t.shape}

What it means

dot_product_attention validates per-dimension shapes of optional operands: -1 means 'any'. This error fires when a fixed expected dimension mismatches, e.g. query's head_dim (last dim) differs from key's, value not matching key shape [B,S,K,H], or mask/bias not 4D-compatible sizes.

Source

Thrown at jax/_src/nn/functions.py:1204

  mask = _ensure_4d(mask) if mask is not None else None
  if query_seq_lengths is not None:
    query_seq_lengths = jnp.asarray(query_seq_lengths)
  if key_value_seq_lengths is not None:
    key_value_seq_lengths = jnp.asarray(key_value_seq_lengths)
  if isinstance(local_window_size, int):
    local_window_size = (local_window_size, local_window_size)

  def _check_shape_and_dtype(t: Array | None, shape: Sequence[int],
                             dtype: DType | None, name: str) -> None:
    if t is None:
      return
    if t.ndim != len(shape):
      raise ValueError(f"{name} ndim should be {len(shape)}, but got {t.ndim}")
    if dtype is not None and t.dtype != dtype:
      raise ValueError(f"{name} dtype should be {dtype}, but got {t.dtype}")
    for i in range(t.ndim):
      if shape[i] != -1 and t.shape[i] != shape[i]:
        raise ValueError(f"{name} shape should be {shape}: but got {t.shape}")

  B, S, K, H = key_arr.shape
  _check_shape_and_dtype(value_arr, [B, S, K, H], key_arr.dtype, 'value')
  _check_shape_and_dtype(query_arr, [B, -1, -1, H], key_arr.dtype, 'query')
  _check_shape_and_dtype(mask, [-1] * 4, np.dtype(bool), 'mask')
  _check_shape_and_dtype(bias, [-1] * 4, None, 'bias')
  _check_shape_and_dtype(query_seq_lengths, [B], np.dtype('int32'),
                         'query_seq_lengths')
  _check_shape_and_dtype(key_value_seq_lengths, [B], np.dtype('int32'),
                         'key_value_seq_lengths')
  if query_arr.shape[-2] % K != 0:
    raise ValueError(f"The number of query heads must be a multiple of "
                     f"key/value heads, but got {query_arr.shape[-2]} vs {K}")

  scale_val = (1.0 / np.sqrt(H)) if scale is None else scale

  match implementation:
    case 'xla':

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make query's last dim equal key's head_dim H and ensure query heads divide/relate to key heads per GQA rules
2. Make value exactly the same shape as key ([B, S, K, H])
3. Verify mask/bias are 4D and broadcast-compatible with [B, S_len, Q_len] scores

Example fix

// before
q = jnp.zeros((8, 128, 16, 128))   # head_dim 128
k = jnp.zeros((8, 128, 16, 64))    # head_dim 64
jax.nn.dot_product_attention(q, k, v)

// after
k = jnp.zeros((8, 128, 16, 128))   # match head_dim
jax.nn.dot_product_attention(q, k, v)

assert q.shape[-1] == k.shape[-1], 'head_dim mismatch'
assert v.shape == k.shape, f'value {v.shape} != key {k.shape}'
assert mask is None or mask.ndim == 4
assert bias is None or bias.ndim == 4

When it happens

Trigger: query with head_dim != key head_dim; value shaped [B,S,H,D] with different S or H from key; mask with shape [B,H,S,S] when broadcast rules of the checker require 4 dims matching -1 pattern (any 4D allowed here since all -1, so practically fires for value/query mismatches).

Common situations: Mismatched model dims between query and key/value projections (GQA misconfiguration); stale checkpoint with changed head_dim; sequence-length mismatch from padding bugs.

Related errors

• {name} ndim should be {len(shape)}, but got {t.ndim}
• The number of query heads must be a multiple of key/value he
• {name} dtype should be {dtype}, but got {t.dtype}
• cuDNN doesn't support right window: {r_window} when causal m
• Unsupported implementation option: {implementation}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/39be6fd5d290cdde. Report an issue: GitHub.