jax-ml/jax · error · ValueError

The dtype of `lengths` must be int32. Got {lengths.dtype}

Error message

The dtype of `lengths` must be int32. Got {lengths.dtype}

What it means

The kernel uses lengths as compile-specialized int32 metadata (compared and used in masking without dtype dispatch), so any other dtype (int64 from numpy defaults, int16, uint32) is rejected. Only exact int32 arrays are accepted.

Source

Thrown at jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py:465

        "Number of Q heads must be divisible by number of KV heads. Got"
        f" {num_q_heads} and {num_kv_heads}."
    )
  if head_dim_k != head_dim:
    raise ValueError(
        "head_dim of Q must be the same as that of K/V. Got"
        f" {head_dim} and {head_dim_k}."
    )
  if pages_per_sequence % pages_per_compute_block != 0:
    raise ValueError(
        "pages_per_compute_block must be divisible by pages per sequence. Got"
        f" {pages_per_compute_block} and {pages_per_sequence}."
    )
  if lengths.shape != (batch_size,):
    raise ValueError("`lengths` and `q` must have the same batch size")
  if batch_size_paged_indices != batch_size:
    raise ValueError("`page_indices` and `q` must have the same batch size")
  if lengths.dtype != jnp.int32:
    raise ValueError(
        f"The dtype of `lengths` must be int32. Got {lengths.dtype}"
    )

  # TODO(dinghua): get the actual cores per chip once there's an official API.
  if megacore_mode == "kv_head":
    if num_kv_heads % 2 != 0:
      raise ValueError(
          "number of KV heads must be even when megacore_mode is 'kv_head'"
      )
    num_cores = 2
  elif megacore_mode == "batch":
    if batch_size % 2 != 0:
      raise ValueError("batch size must be even when megacore_mode is 'batch'")
    num_cores = 2
  elif megacore_mode is None:
    num_cores = 1
  else:
    raise ValueError("megacore_mode must be one of ['kv_head', 'batch', None]")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert explicitly: lengths = jnp.asarray(lengths, dtype=jnp.int32)
2. Build lengths with dtype=jnp.int32 at creation (np.array(x, dtype=np.int32))
3. Audit dataloader output dtypes; TPU Pallas kernels commonly require int32 indices

Example fix

// before
lengths = np.array([12, 34, 56])  # int64
// after
lengths = np.array([12, 34, 56], dtype=np.int32)

lengths = jnp.asarray(lengths, dtype=jnp.int32)

def is_int32(a): return isinstance(a, jax.Array) and a.dtype == jnp.int32

When it happens

Trigger: Passing lengths produced by np.arange(...)/np.array(...) (default int64 on Linux), jnp.array(..., jnp.int64), or values from a tokenizer that returns int64; also Python lists converted lazily.

Common situations: Numpy's platform default int64 leaking into JAX TPU code; JAX config jax_enable_x64=True making integer literals int64; loading lengths from a dataset stored as int64.

Related errors

• The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must
• k_pages and v_pages must have the same shape. Got {k_pages.s
• Number of Q heads must be divisible by number of KV heads. G
• head_dim of Q must be the same as that of K/V. Got {head_dim
• pages_per_compute_block must be divisible by pages per seque

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/777e324ad2318eef. Report an issue: GitHub.