jax-ml/jax · error · ValueError

{name} ndim should be {len(shape)}, but got {t.ndim}

Error message

{name} ndim should be {len(shape)}, but got {t.ndim}

What it means

jax.nn.dot_product_attention validates each optional operand (query, value, mask, bias, seq lengths) against an expected rank. This error fires when an operand's ndim does not match the required number of dimensions for its role (e.g. query/key/value must be 4D BSHD; mask/bias 4D).

Source

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

  query_arr = _ensure_4d(query)
  key_arr = _ensure_4d(key)
  value_arr = _ensure_4d(value)
  bias = _ensure_4d(bias) if bias is not None else None
  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}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape/rearrange inputs to 4D [batch, seq_len, num_heads, head_dim] (e.g. with einops.rearrange 'b h s d -> b s h d')
2. Add a batch dimension with x[None] for single-sequence inputs
3. Check the docstring of jax.nn.dot_product_attention for the exact operand layouts

Example fix

// before
out = jax.nn.dot_product_attention(q3, k3, v3)  # (B, H, S, D)

// after
import jax.numpy as jnp
q, k, v = (jnp.transpose(a, (0, 2, 1, 3)) for a in (q3, k3, v3))  # BSHD
out = jax.nn.dot_product_attention(q, k, v)

def check_4d(name, t):
    if t is not None and t.ndim != 4:
        raise ValueError(f'{name} must be 4D (B,S,H,D), got ndim={t.ndim}')
check_4d('query', q); check_4d('key', k); check_4d('value', v)
out = jax.nn.dot_product_attention(q, k, v)

def is_bshd(t) -> bool: return getattr(t, 'ndim', 0) == 4

When it happens

Trigger: Passing a 3D (BHSD-layout) or 2D attention input to dot_product_attention which expects 4D [B, S, H, D]; passing a 1D mask; passing query_seq_lengths with ndim != 1.

Common situations: Porting code written for BHSD-layout attention (transformer implementations, Flax attention) into jax.nn.dot_product_attention; forgetting the leading batch dim; passing packed 2D sequences.

Related errors

• {name} shape should be {shape}: but got {t.shape}
• 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/1bcca8b899f00f8d. Report an issue: GitHub.