{"record":{"id":"777e324ad2318eef","repo":"jax-ml/jax","slug":"the-dtype-of-lengths-must-be-int32-got-lengths-777e32","errorCode":null,"errorMessage":"The dtype of `lengths` must be int32. Got {lengths.dtype}","messagePattern":"The dtype of `lengths` must be int32\\. Got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py","lineNumber":465,"sourceCode":"        \"Number of Q heads must be divisible by number of KV heads. Got\"\n        f\" {num_q_heads} and {num_kv_heads}.\"\n    )\n  if head_dim_k != head_dim:\n    raise ValueError(\n        \"head_dim of Q must be the same as that of K/V. Got\"\n        f\" {head_dim} and {head_dim_k}.\"\n    )\n  if pages_per_sequence % pages_per_compute_block != 0:\n    raise ValueError(\n        \"pages_per_compute_block must be divisible by pages per sequence. Got\"\n        f\" {pages_per_compute_block} and {pages_per_sequence}.\"\n    )\n  if lengths.shape != (batch_size,):\n    raise ValueError(\"`lengths` and `q` must have the same batch size\")\n  if batch_size_paged_indices != batch_size:\n    raise ValueError(\"`page_indices` and `q` must have the same batch size\")\n  if lengths.dtype != jnp.int32:\n    raise ValueError(\n        f\"The dtype of `lengths` must be int32. Got {lengths.dtype}\"\n    )\n\n  # TODO(dinghua): get the actual cores per chip once there's an official API.\n  if megacore_mode == \"kv_head\":\n    if num_kv_heads % 2 != 0:\n      raise ValueError(\n          \"number of KV heads must be even when megacore_mode is 'kv_head'\"\n      )\n    num_cores = 2\n  elif megacore_mode == \"batch\":\n    if batch_size % 2 != 0:\n      raise ValueError(\"batch size must be even when megacore_mode is 'batch'\")\n    num_cores = 2\n  elif megacore_mode is None:\n    num_cores = 1\n  else:\n    raise ValueError(\"megacore_mode must be one of ['kv_head', 'batch', None]\")","sourceCodeStart":447,"sourceCodeEnd":483,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py#L447-L483","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Convert explicitly: lengths = jnp.asarray(lengths, dtype=jnp.int32)","Build lengths with dtype=jnp.int32 at creation (np.array(x, dtype=np.int32))","Audit dataloader output dtypes; TPU Pallas kernels commonly require int32 indices"],"exampleFix":"// before\nlengths = np.array([12, 34, 56])  # int64\n// after\nlengths = np.array([12, 34, 56], dtype=np.int32)","handlingStrategy":"type-guard","validationCode":"lengths = jnp.asarray(lengths, dtype=jnp.int32)","typeGuard":"def is_int32(a): return isinstance(a, jax.Array) and a.dtype == jnp.int32","tryCatchPattern":null,"preventionTips":["Coerce all index-like tensors to int32 at the dataloader boundary","Watch out for numpy int64 defaults and jax_enable_x64=True on TPU"],"tags":["jax","pallas","tpu","paged-attention","dtype-validation"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}