jax-ml/jax · error · ValueError

{name} dtype should be {dtype}, but got {t.dtype}

Error message

{name} dtype should be {dtype}, but got {t.dtype}

What it means

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.

Source

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

  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}")

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast operands consistently: q = q.astype(k.dtype), v = v.astype(k.dtype)
2. Convert mask with mask.astype(bool) and seq lengths with .astype(jnp.int32)
3. Set inputs in one dtype up front (e.g. jnp.bfloat16) before calling attention

Example fix

// before
out = jax.nn.dot_product_attention(q_f32, k_bf16, v_bf16)

// after
out = jax.nn.dot_product_attention(q_f32.astype(k_bf16.dtype), k_bf16, v_bf16)

target = k.dtype
q, v = q.astype(target), v.astype(target)
mask = None if mask is None else mask.astype(bool)
qsl = None if qsl is None else qsl.astype(jnp.int32)

When it happens

Trigger: 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.

Common situations: Mixed-precision training where only some attention projections were cast; loading checkpoints with different default dtypes; numpy interop producing int64 lengths on Linux.

Related errors

• {name} ndim should be {len(shape)}, but got {t.ndim}
• {name} shape should be {shape}: but got {t.shape}
• The number of query heads must be a multiple of key/value he
• 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/32c5cdfbeeafb325. Report an issue: GitHub.