{"record":{"id":"f5e0f6fa5e8f0748","repo":"jax-ml/jax","slug":"masked-swap-with-non-32-bit-data","errorCode":null,"errorMessage":"masked swap with non-32-bit data","messagePattern":"masked swap with non-32-bit data","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2480,"sourceCode":"  int_out_type = ctx.aval_to_ir_type(expected_aval, is_kernel_boundary=True)\n  return arith.extui(int_out_type, val)\n\n\n@register_lowering_rule(primitives.swap_p, ensure_mlir_values=False)\ndef _masked_swap_lowering_rule(\n    ctx: LoweringRuleContext, *args_flat, args_tree, **_\n):\n  ref, transforms, val, mask = args_tree.unflatten(args_flat)\n  ref_aval, transforms_avals, val_aval, mask_aval = args_tree.unflatten(\n      ctx.avals_in\n  )\n  prev_transforms, idx = _canonicalize_transforms_to_indexer(\n      ref_aval, transforms, transforms_avals\n  )\n\n  if mask is not None:\n    if  val_aval.dtype.itemsize != 4:\n      raise NotImplementedError(\"masked swap with non-32-bit data\")\n    if val_aval.shape != mask_aval.shape:\n      raise ValueError(\n          \"Expected value and mask to have the same shape, but got\"\n          f\" value shape {val_aval.shape} vs. mask shape {mask_aval.shape}.\"\n      )\n\n  ref_block_shape, *_ = ctx.block_shapes\n  ref, ref_block_shape = _transform_ref(\n      ref, ref_aval, ref_block_shape, prev_transforms\n  )\n\n  ref_type = ir.MemRefType(ref.type)\n  memory_space = str(ref_type.memory_space)\n  is_smem_store = memory_space == \"#tpu.memory_space<smem>\"\n  is_vmem_store = memory_space == \"#tpu.memory_space<vmem>\"\n  (aval_out,) = ctx.avals_out\n  if not isinstance(val, ir.Value):\n    val = ir_constant(val, mlir_type=_dtype_to_ir_type(val_aval.dtype))","sourceCodeStart":2462,"sourceCodeEnd":2498,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2462-L2498","documentation":"Raised when a masked swap (pl.store with a mask) on TPU operates on data whose element size is not 32 bits (4 bytes). Masked vector stores are only implemented for 32-bit lanes (f32/i32); 8/16/64-bit dtypes with masks are unsupported.","triggerScenarios":"pl.store(value, ref, mask=mask) where value dtype is e.g. f16, bf16, int8, f64 — any itemsize != 4 — inside a TPU Pallas kernel.","commonSituations":"Writing half-precision kernels with predicated stores; using boolean/int8 outputs with masks for sparsity.","solutions":["Store 32-bit data: cast values to jnp.float32/int32 before the masked store (and cast back on the next load)","Drop the mask: compute the full block and store unmasked, relying on block padding","Pad the value to 32 bits manually and bitcast"],"exampleFix":"// before\npl.store(ref, x_f16, mask=m)\n// after\npl.store(ref, x_f16.astype(jnp.float32), mask=m)","handlingStrategy":"fallback","validationCode":"def masked_store(ref, val, mask):\n    if val.dtype.itemsize != 4:\n        val = val.astype(jnp.float32)\n    pl.store(ref, val, mask=mask)","typeGuard":"def mask_store_supported(dtype) -> bool:\n    import numpy as np\n    return np.dtype(dtype).itemsize == 4","tryCatchPattern":null,"preventionTips":["Standardize kernels on f32/i32 for masked outputs","Check dtype.itemsize == 4 before any masked store"],"tags":["jax","pallas","tpu","masked-store","dtype"],"backgroundTag":"pallas-masked-store-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}