{"record":{"id":"39be6fd5d290cdde","repo":"jax-ml/jax","slug":"name-shape-should-be-shape-but-got-t-shape","errorCode":null,"errorMessage":"{name} shape should be {shape}: but got {t.shape}","messagePattern":"(.+?) shape should be (.+?): but got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/nn/functions.py","lineNumber":1204,"sourceCode":"  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}\")\n\n  scale_val = (1.0 / np.sqrt(H)) if scale is None else scale\n\n  match implementation:\n    case 'xla':","sourceCodeStart":1186,"sourceCodeEnd":1222,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/nn/functions.py#L1186-L1222","documentation":"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.","triggerScenarios":"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).","commonSituations":"Mismatched model dims between query and key/value projections (GQA misconfiguration); stale checkpoint with changed head_dim; sequence-length mismatch from padding bugs.","solutions":["Make query's last dim equal key's head_dim H and ensure query heads divide/relate to key heads per GQA rules","Make value exactly the same shape as key ([B, S, K, H])","Verify mask/bias are 4D and broadcast-compatible with [B, S_len, Q_len] scores"],"exampleFix":"// before\nq = jnp.zeros((8, 128, 16, 128))   # head_dim 128\nk = jnp.zeros((8, 128, 16, 64))    # head_dim 64\njax.nn.dot_product_attention(q, k, v)\n\n// after\nk = jnp.zeros((8, 128, 16, 128))   # match head_dim\njax.nn.dot_product_attention(q, k, v)","handlingStrategy":"validation","validationCode":"assert q.shape[-1] == k.shape[-1], 'head_dim mismatch'\nassert v.shape == k.shape, f'value {v.shape} != key {k.shape}'\nassert mask is None or mask.ndim == 4\nassert bias is None or bias.ndim == 4","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Validate shapes once in a setup/pre-flight function, not inside jit","Log operand shapes on setup to catch config drift early"],"tags":["jax","nn","attention","shape-validation"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}