jax-ml/jax · error · ValueError

Mask function must return a boolean-valued array, but got: {

Error message

Mask function must return a boolean-valued array, but got: {computed_mask.dtype}

What it means

The custom mask function passed to the Splash Attention TPU kernel must return an array of dtype jnp.bool_, but the user-supplied callable returned a different dtype (e.g. int32 or bfloat16 after comparisons/arithmetic). The kernel validates the dtype because it uses the mask directly as a boolean multiplier inside the Pallas kernel, and a non-boolean dtype would produce wrong results or fail to compile.

Source

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

      repeats, rem = divmod(k_slice.size, NUM_LANES)
      assert rem == 0
      q_sequence = jnp.tile(
          q_sequence_ref[...], (1, repeats)
      )  # [bq, k_slice.size]
    else:
      assert q_sequence_ref.shape == (NUM_SUBLANES, bq)

      k_sequence = k_offset + jax.lax.broadcasted_iota(
          jnp.int32, (k_slice.size, bq), 0
      )
      q_sequence = q_sequence_ref[:1, :]  # [1, bq]
      q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq))

    assert q_sequence.shape == k_sequence.shape
    computed_mask = mask_function(q_sequence, k_sequence)
    if computed_mask.dtype != jnp.dtype(jnp.bool_):
      raise ValueError(
          "Mask function must return a boolean-valued array, but got:"
          f" {computed_mask.dtype}"
      )
    masks.append(computed_mask)

  if q_segment_ids_ref is not None:
    if k_in_lanes:
      kv_ids = kv_segment_ids_ref[:1, k_slice]  # [1, k_slice]
      repeats, rem = divmod(kv_ids.shape[1], NUM_LANES)
      if rem:
        raise NotImplementedError(f"block_kv must be a multiple of {NUM_LANES}")
      q_ids = jnp.tile(q_segment_ids_ref[:], (1, repeats))  # [bq, bkv]
    else:
      assert bq == q_segment_ids_ref.shape[-1]
      repeats, rem = divmod(bq, NUM_LANES)
      if rem:
        raise NotImplementedError(f"block_q must be a multiple of {NUM_LANES}")
      kv_ids = jnp.tile(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Return an explicit boolean from the mask function: wrap the expression in .astype(jnp.bool_) or jnp.asarray(..., dtype=jnp.bool_)
2. Verify the mask function signature matches (q_sequence, k_sequence) -> bool array and contains only comparison/logical ops (>, >=, ==, &, |, ~)
3. Use the built-in mask helpers from splash_attention_mask instead of a custom function

Example fix

// before
mask_function=lambda q, k: (q[:, None] >= k[None, :]) * 1.0
// after
mask_function=lambda q, k: (q >= k).astype(jnp.bool_)

def check_mask_fn(mask_fn, q_len, kv_len):
    out = mask_fn(jnp.arange(q_len)[:, None], jnp.arange(kv_len)[None, :])
    assert out.dtype == jnp.bool_, f'mask must be bool, got {out.dtype}'
    return out.shape == (q_len, kv_len)

def is_bool_mask_fn(fn) -> bool:
    out = fn(jnp.zeros((4, 4), jnp.int32), jnp.zeros((4, 4), jnp.int32))
    return jnp.dtype(out.dtype) == jnp.dtype(jnp.bool_)

When it happens

Trigger: Calling make_splash_attention / splash_attention_kernel with a mask_function argument whose body returns e.g. q_sequence[:, None] > k_sequence (fine) combined with arithmetic like (q < k) * 1, or jnp.where(...) defaulting to int, or returning a float score instead of a boolean.

Common situations: Porting a mask from another attention implementation where masks were float (0./-inf additive masks); writing mask_function as lambda q, k: (q >= k).astype(jnp.int32); using jnp.where which promotes dtype.

Related errors

• partial_mask_blocks must be of type np.bool_ but got {partia
• 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/b27c9fc572f8239e. Report an issue: GitHub.