jax-ml/jax · error · NotImplementedError

masked swap with non-32-bit data

Error message

masked swap with non-32-bit data

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:2480

  int_out_type = ctx.aval_to_ir_type(expected_aval, is_kernel_boundary=True)
  return arith.extui(int_out_type, val)


@register_lowering_rule(primitives.swap_p, ensure_mlir_values=False)
def _masked_swap_lowering_rule(
    ctx: LoweringRuleContext, *args_flat, args_tree, **_
):
  ref, transforms, val, mask = args_tree.unflatten(args_flat)
  ref_aval, transforms_avals, val_aval, mask_aval = args_tree.unflatten(
      ctx.avals_in
  )
  prev_transforms, idx = _canonicalize_transforms_to_indexer(
      ref_aval, transforms, transforms_avals
  )

  if mask is not None:
    if  val_aval.dtype.itemsize != 4:
      raise NotImplementedError("masked swap with non-32-bit data")
    if val_aval.shape != mask_aval.shape:
      raise ValueError(
          "Expected value and mask to have the same shape, but got"
          f" value shape {val_aval.shape} vs. mask shape {mask_aval.shape}."
      )

  ref_block_shape, *_ = ctx.block_shapes
  ref, ref_block_shape = _transform_ref(
      ref, ref_aval, ref_block_shape, prev_transforms
  )

  ref_type = ir.MemRefType(ref.type)
  memory_space = str(ref_type.memory_space)
  is_smem_store = memory_space == "#tpu.memory_space<smem>"
  is_vmem_store = memory_space == "#tpu.memory_space<vmem>"
  (aval_out,) = ctx.avals_out
  if not isinstance(val, ir.Value):
    val = ir_constant(val, mlir_type=_dtype_to_ir_type(val_aval.dtype))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Store 32-bit data: cast values to jnp.float32/int32 before the masked store (and cast back on the next load)
2. Drop the mask: compute the full block and store unmasked, relying on block padding
3. Pad the value to 32 bits manually and bitcast

Example fix

// before
pl.store(ref, x_f16, mask=m)
// after
pl.store(ref, x_f16.astype(jnp.float32), mask=m)

def masked_store(ref, val, mask):
    if val.dtype.itemsize != 4:
        val = val.astype(jnp.float32)
    pl.store(ref, val, mask=mask)

def mask_store_supported(dtype) -> bool:
    import numpy as np
    return np.dtype(dtype).itemsize == 4

When it happens

Trigger: pl.store(value, ref, mask=mask) where value dtype is e.g. f16, bf16, int8, f64 — any itemsize != 4 — inside a TPU Pallas kernel.

Common situations: Writing half-precision kernels with predicated stores; using boolean/int8 outputs with masks for sparsity.

Related errors

• Acc ref dtype must be float32 or int32, got {dtype}
• Expected value and mask to have the same shape, but got valu
• SMEM store does not support masks
• masked swap with strided store
• Reductions over unsigned integers not implemented.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/f5e0f6fa5e8f0748. Report an issue: GitHub.