{"record":{"id":"1bcca8b899f00f8d","repo":"jax-ml/jax","slug":"name-ndim-should-be-len-shape-but-got-t-ndi","errorCode":null,"errorMessage":"{name} ndim should be {len(shape)}, but got {t.ndim}","messagePattern":"(.+?) ndim should be (.+?), but got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/nn/functions.py","lineNumber":1199,"sourceCode":"\n  query_arr = _ensure_4d(query)\n  key_arr = _ensure_4d(key)\n  value_arr = _ensure_4d(value)\n  bias = _ensure_4d(bias) if bias is not None else None\n  mask = _ensure_4d(mask) if mask is not None else None\n  if query_seq_lengths is not None:\n    query_seq_lengths = jnp.asarray(query_seq_lengths)\n  if key_value_seq_lengths is not None:\n    key_value_seq_lengths = jnp.asarray(key_value_seq_lengths)\n  if isinstance(local_window_size, int):\n    local_window_size = (local_window_size, local_window_size)\n\n  def _check_shape_and_dtype(t: Array | None, shape: Sequence[int],\n                             dtype: DType | None, name: str) -> None:\n    if t is None:\n      return\n    if t.ndim != len(shape):\n      raise ValueError(f\"{name} ndim should be {len(shape)}, but got {t.ndim}\")\n    if dtype is not None and t.dtype != dtype:\n      raise ValueError(f\"{name} dtype should be {dtype}, but got {t.dtype}\")\n    for i in range(t.ndim):\n      if shape[i] != -1 and t.shape[i] != shape[i]:\n        raise ValueError(f\"{name} shape should be {shape}: but got {t.shape}\")\n\n  B, S, K, H = key_arr.shape\n  _check_shape_and_dtype(value_arr, [B, S, K, H], key_arr.dtype, 'value')\n  _check_shape_and_dtype(query_arr, [B, -1, -1, H], key_arr.dtype, 'query')\n  _check_shape_and_dtype(mask, [-1] * 4, np.dtype(bool), 'mask')\n  _check_shape_and_dtype(bias, [-1] * 4, None, 'bias')\n  _check_shape_and_dtype(query_seq_lengths, [B], np.dtype('int32'),\n                         'query_seq_lengths')\n  _check_shape_and_dtype(key_value_seq_lengths, [B], np.dtype('int32'),\n                         'key_value_seq_lengths')\n  if query_arr.shape[-2] % K != 0:\n    raise ValueError(f\"The number of query heads must be a multiple of \"\n                     f\"key/value heads, but got {query_arr.shape[-2]} vs {K}\")","sourceCodeStart":1181,"sourceCodeEnd":1217,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/nn/functions.py#L1181-L1217","documentation":"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).","triggerScenarios":"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.","commonSituations":"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.","solutions":["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')","Add a batch dimension with x[None] for single-sequence inputs","Check the docstring of jax.nn.dot_product_attention for the exact operand layouts"],"exampleFix":"// before\nout = jax.nn.dot_product_attention(q3, k3, v3)  # (B, H, S, D)\n\n// after\nimport jax.numpy as jnp\nq, k, v = (jnp.transpose(a, (0, 2, 1, 3)) for a in (q3, k3, v3))  # BSHD\nout = jax.nn.dot_product_attention(q, k, v)","handlingStrategy":"validation","validationCode":"def check_4d(name, t):\n    if t is not None and t.ndim != 4:\n        raise ValueError(f'{name} must be 4D (B,S,H,D), got ndim={t.ndim}')\ncheck_4d('query', q); check_4d('key', k); check_4d('value', v)\nout = jax.nn.dot_product_attention(q, k, v)","typeGuard":"def is_bshd(t) -> bool: return getattr(t, 'ndim', 0) == 4","tryCatchPattern":null,"preventionTips":["Standardize on BSHD layout across the model; add rank asserts at module boundaries","Write a single _to_bshd(x) helper used everywhere before attention"],"tags":["jax","nn","attention","shape-validation"],"backgroundTag":"rank-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}