{"record":{"id":"b27c9fc572f8239e","repo":"jax-ml/jax","slug":"mask-function-must-return-a-boolean-valued-array","errorCode":null,"errorMessage":"Mask function must return a boolean-valued array, but got: {computed_mask.dtype}","messagePattern":"Mask function must return a boolean-valued array, but got: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":657,"sourceCode":"\n      repeats, rem = divmod(k_slice.size, NUM_LANES)\n      assert rem == 0\n      q_sequence = jnp.tile(\n          q_sequence_ref[...], (1, repeats)\n      )  # [bq, k_slice.size]\n    else:\n      assert q_sequence_ref.shape == (NUM_SUBLANES, bq)\n\n      k_sequence = k_offset + jax.lax.broadcasted_iota(\n          jnp.int32, (k_slice.size, bq), 0\n      )\n      q_sequence = q_sequence_ref[:1, :]  # [1, bq]\n      q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq))\n\n    assert q_sequence.shape == k_sequence.shape\n    computed_mask = mask_function(q_sequence, k_sequence)\n    if computed_mask.dtype != jnp.dtype(jnp.bool_):\n      raise ValueError(\n          \"Mask function must return a boolean-valued array, but got:\"\n          f\" {computed_mask.dtype}\"\n      )\n    masks.append(computed_mask)\n\n  if q_segment_ids_ref is not None:\n    if k_in_lanes:\n      kv_ids = kv_segment_ids_ref[:1, k_slice]  # [1, k_slice]\n      repeats, rem = divmod(kv_ids.shape[1], NUM_LANES)\n      if rem:\n        raise NotImplementedError(f\"block_kv must be a multiple of {NUM_LANES}\")\n      q_ids = jnp.tile(q_segment_ids_ref[:], (1, repeats))  # [bq, bkv]\n    else:\n      assert bq == q_segment_ids_ref.shape[-1]\n      repeats, rem = divmod(bq, NUM_LANES)\n      if rem:\n        raise NotImplementedError(f\"block_q must be a multiple of {NUM_LANES}\")\n      kv_ids = jnp.tile(","sourceCodeStart":639,"sourceCodeEnd":675,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L639-L675","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Return an explicit boolean from the mask function: wrap the expression in .astype(jnp.bool_) or jnp.asarray(..., dtype=jnp.bool_)","Verify the mask function signature matches (q_sequence, k_sequence) -> bool array and contains only comparison/logical ops (>, >=, ==, &, |, ~)","Use the built-in mask helpers from splash_attention_mask instead of a custom function"],"exampleFix":"// before\nmask_function=lambda q, k: (q[:, None] >= k[None, :]) * 1.0\n// after\nmask_function=lambda q, k: (q >= k).astype(jnp.bool_)","handlingStrategy":"validation","validationCode":"def check_mask_fn(mask_fn, q_len, kv_len):\n    out = mask_fn(jnp.arange(q_len)[:, None], jnp.arange(kv_len)[None, :])\n    assert out.dtype == jnp.bool_, f'mask must be bool, got {out.dtype}'\n    return out.shape == (q_len, kv_len)","typeGuard":"def is_bool_mask_fn(fn) -> bool:\n    out = fn(jnp.zeros((4, 4), jnp.int32), jnp.zeros((4, 4), jnp.int32))\n    return jnp.dtype(out.dtype) == jnp.dtype(jnp.bool_)","tryCatchPattern":null,"preventionTips":["Always end custom mask functions with .astype(jnp.bool_)","Prefer library-provided mask helpers over hand-written masks","Unit-test the mask function standalone before passing to the kernel"],"tags":["jax","pallas","tpu","splash-attention","dtype","mask"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}