jax-ml/jax · error · ValueError

partial_mask_blocks must be of type np.bool_ but got {partia

Error message

partial_mask_blocks must be of type np.bool_ but got {partial_mask_blocks.dtype}

What it means

The optional partial_mask_blocks array supplied via the mask info (used for sparse/partial block masking in Splash Attention) must be a numpy bool array. Passing any other dtype (int, uint8, float) is rejected because the kernel indexes grid blocks with it directly.

Source

Thrown at jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py:938

  bkv_compute = block_sizes.block_kv_compute

  if is_mqa:
    expected_kv_rank = 2
    kv_head_dimension = 1
    kv_seq_len_dimension = 0
    num_kv_heads = 1
  else:
    expected_kv_rank = 3
    kv_head_dimension = 2
    kv_seq_len_dimension = 1
    num_kv_heads = k.shape[0]

  partial_mask_blocks = fwd_mask_info.partial_mask_blocks
  if (
      partial_mask_blocks is not None
      and jnp.dtype(partial_mask_blocks.dtype) != np.bool_
  ):
    raise ValueError(
        "partial_mask_blocks must be of type np.bool_ but got"
        f" {partial_mask_blocks.dtype}"
    )

  if len(k.shape) != expected_kv_rank:
    raise ValueError(
        f"Expected {expected_kv_rank}-dim 'key' tensor for MQA. Instead got a"
        f" {len(k.shape)}-dim one."
    )

  if k.shape[kv_head_dimension] != head_dim_qk:
    raise ValueError(
        f"Expected 'key' head dimension to be: {head_dim_qk}. Instead got:"
        f" {k.shape[kv_head_dimension]}."
    )

  if not is_mqa and num_q_heads % num_kv_heads != 0:
    raise ValueError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the array before passing: partial_mask_blocks = partial_mask_blocks.astype(np.bool_)
2. Keep the output of mask-computation helpers (which already return bool) instead of round-tripping through other dtypes
3. If memory was the reason for uint8, keep the mask in numpy (not jnp) bool to avoid device copies

Example fix

// before
partial = (scores > 0).astype(np.uint8)
// after
partial = (scores > 0).astype(np.bool_)

partial = np.asarray(partial, dtype=np.bool_)

def is_bool_mask(a) -> bool:
    return isinstance(a, np.ndarray) and a.dtype == np.bool_

When it happens

Trigger: Calling make_splash_attention_attention_function/make_fast_softmax with a MultiBlockMask with partial_mask_blocks computed as np.arange(...) % 2 or cast to uint8 to save memory; also passing a jnp.uint8 array created with jnp.bool_ conversions.

Common situations: Building custom BlockMask sparsity patterns and casting masks to uint8/int8 to reduce transfer size; converting masks with .view(np.uint8).

Related errors

• Mask function must return a boolean-valued array, but got: {
• Acc ref dtype must be float32 or int32, got {dtype}
• Out-of-bounds masked swap of ({device_id} {local_core_id} {m
• 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/57a44e9493035600. Report an issue: GitHub.