{"record":{"id":"32c5cdfbeeafb325","repo":"jax-ml/jax","slug":"name-dtype-should-be-dtype-but-got-t-dtype","errorCode":null,"errorMessage":"{name} dtype should be {dtype}, but got {t.dtype}","messagePattern":"(.+?) dtype should be (.+?), but got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/nn/functions.py","lineNumber":1201,"sourceCode":"  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}\")\n\n  scale_val = (1.0 / np.sqrt(H)) if scale is None else scale","sourceCodeStart":1183,"sourceCodeEnd":1219,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/nn/functions.py#L1183-L1219","documentation":"dot_product_attention checks optional operands against the dtype of the key tensor (or bool for masks, int32 for sequence lengths). This error fires when an operand has the right shape but a different dtype, e.g. query in float32 while key is bfloat16.","triggerScenarios":"Passing query/value arrays cast to a different precision than key (mixed float32/bfloat16 inputs); mask as uint8 instead of bool; query_seq_lengths as int64 instead of int32.","commonSituations":"Mixed-precision training where only some attention projections were cast; loading checkpoints with different default dtypes; numpy interop producing int64 lengths on Linux.","solutions":["Cast operands consistently: q = q.astype(k.dtype), v = v.astype(k.dtype)","Convert mask with mask.astype(bool) and seq lengths with .astype(jnp.int32)","Set inputs in one dtype up front (e.g. jnp.bfloat16) before calling attention"],"exampleFix":"// before\nout = jax.nn.dot_product_attention(q_f32, k_bf16, v_bf16)\n\n// after\nout = jax.nn.dot_product_attention(q_f32.astype(k_bf16.dtype), k_bf16, v_bf16)","handlingStrategy":"validation","validationCode":"target = k.dtype\nq, v = q.astype(target), v.astype(target)\nmask = None if mask is None else mask.astype(bool)\nqsl = None if qsl is None else qsl.astype(jnp.int32)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Cast all attention operands once at module entry","Add dtype asserts in mixed-precision training loops"],"tags":["jax","nn","attention","dtype-validation"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}