{"record":{"id":"57a44e9493035600","repo":"jax-ml/jax","slug":"partial-mask-blocks-must-be-of-type-np-bool-but-g","errorCode":null,"errorMessage":"partial_mask_blocks must be of type np.bool_ but got {partial_mask_blocks.dtype}","messagePattern":"partial_mask_blocks must be of type np\\.bool_ but got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":938,"sourceCode":"  bkv_compute = block_sizes.block_kv_compute\n\n  if is_mqa:\n    expected_kv_rank = 2\n    kv_head_dimension = 1\n    kv_seq_len_dimension = 0\n    num_kv_heads = 1\n  else:\n    expected_kv_rank = 3\n    kv_head_dimension = 2\n    kv_seq_len_dimension = 1\n    num_kv_heads = k.shape[0]\n\n  partial_mask_blocks = fwd_mask_info.partial_mask_blocks\n  if (\n      partial_mask_blocks is not None\n      and jnp.dtype(partial_mask_blocks.dtype) != np.bool_\n  ):\n    raise ValueError(\n        \"partial_mask_blocks must be of type np.bool_ but got\"\n        f\" {partial_mask_blocks.dtype}\"\n    )\n\n  if len(k.shape) != expected_kv_rank:\n    raise ValueError(\n        f\"Expected {expected_kv_rank}-dim 'key' tensor for MQA. Instead got a\"\n        f\" {len(k.shape)}-dim one.\"\n    )\n\n  if k.shape[kv_head_dimension] != head_dim_qk:\n    raise ValueError(\n        f\"Expected 'key' head dimension to be: {head_dim_qk}. Instead got:\"\n        f\" {k.shape[kv_head_dimension]}.\"\n    )\n\n  if not is_mqa and num_q_heads % num_kv_heads != 0:\n    raise ValueError(","sourceCodeStart":920,"sourceCodeEnd":956,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L920-L956","documentation":"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.","triggerScenarios":"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.","commonSituations":"Building custom BlockMask sparsity patterns and casting masks to uint8/int8 to reduce transfer size; converting masks with .view(np.uint8).","solutions":["Cast the array before passing: partial_mask_blocks = partial_mask_blocks.astype(np.bool_)","Keep the output of mask-computation helpers (which already return bool) instead of round-tripping through other dtypes","If memory was the reason for uint8, keep the mask in numpy (not jnp) bool to avoid device copies"],"exampleFix":"// before\npartial = (scores > 0).astype(np.uint8)\n// after\npartial = (scores > 0).astype(np.bool_)","handlingStrategy":"type-guard","validationCode":"partial = np.asarray(partial, dtype=np.bool_)","typeGuard":"def is_bool_mask(a) -> bool:\n    return isinstance(a, np.ndarray) and a.dtype == np.bool_","tryCatchPattern":null,"preventionTips":["Never store masks as uint8/int8 for this API","Cast with .astype(np.bool_) at the boundary"],"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"}