jax-ml/jax · error · ValueError

Expected int32 input, but got {array.dtype}.

Error message

Expected int32 input, but got {array.dtype}.

What it means

The splash attention mask preprocessing downcasts mask index arrays from int32 to int8/int16 to save TPU memory. _downcast only accepts int32 arrays; any other integer width (int64, uint32, etc.) is rejected before the astype.

Source

Thrown at jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py:482

    data_next_per_head_list.append(data_next_per_head)
    mask_next_per_head = jnp.concatenate(
        mask_next_sequence_slices, axis=q_sequence_axis
    )
    mask_next_per_head_list.append(mask_next_per_head)

  # Concatenate (or broadcast) the head shards.
  data_next = jnp.concatenate(data_next_per_head_list, axis=head_axis)
  mask_next = jnp.concatenate(mask_next_per_head_list, axis=head_axis)

  if is_dkv:
    partial_mask_blocks = partial_mask_blocks.swapaxes(-1, -2)

  def _downcast(array: jax.Array, max_value: int) -> jax.Array:
    if array.size == 0:
      return array

    if array.dtype != np.int32:
      raise ValueError(f'Expected int32 input, but got {array.dtype}.')

    if max_value <= np.iinfo(np.int8).max:
      return array.astype(np.int8)
    elif max_value <= np.iinfo(np.int16).max:
      return array.astype(np.int16)
    else:
      return array.astype(np.int32)

  if downcast_smem_data:
    block_mask = block_mask.astype(np.int8)  # values are in the range [0, 1, 2]
    data_next = _downcast(
        data_next, q_blocks_per_shard if is_dkv else kv_blocks_count
    )
    mask_next = _downcast(
        mask_next, heads_per_shard * q_blocks_per_shard * kv_blocks_count
    )

  return (

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast mask index arrays to jnp.int32 before passing them: arr.astype(jnp.int32)
2. Audit mask construction under jax_enable_x64=True; wrap dynamic mask creation in a helper that forces int32

Example fix

# before
dynamic_mask = (positions.astype(jnp.int64))  # x64 enabled
# after
dynamic_mask = (positions.astype(jnp.int32))

assert dynamic_mask_indices.dtype == jnp.int32

def is_int32(a) -> bool:
    return a.dtype == jnp.int32

When it happens

Trigger: Passing a dynamic mask index array created with jnp.arange(..., dtype=jnp.int64) or numpy int64/uint arrays to make_splash_attention_mask's dynamic-mask path. On 64-bit-enabled JAX (jax_enable_x64=True) literals and arange default to int64 and trigger this.

Common situations: Enabling jax_enable_x64 in a training script then reusing the same mask-building code; constructing mask indices with numpy defaults (int64 on Linux) instead of jnp.int32.

Related errors

• Mask function must return a boolean-valued array, but got: {
• partial_mask_blocks must be of type np.bool_ but got {partia
• Acc ref dtype must be float32 or int32, got {dtype}
• masked swap with non-32-bit data
• Reductions over unsigned integers not implemented.

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