jax-ml/jax · error · TypeError

mask.dtype={mask.dtype} is not a boolean dtype

Error message

mask.dtype={mask.dtype} is not a boolean dtype

What it means

scan_count requires mask to have a boolean dtype; integer or float masks are rejected with a TypeError.

Source

Thrown at jax/_src/pallas/mosaic/sc_primitives.py:529

  """Blocks until all subcores on the same core reach this instruction.

  The barrier must be used with
  :class:`jax.experimental.pallas.tpu_sc.VectorSubcoreMesh`.
  """
  barrier_p.bind()


scan_count_p = jax_core.Primitive("scan_count")
scan_count_p.multiple_results = True


@scan_count_p.def_abstract_eval
def _scan_count_abstract_eval(x, mask):
  if x.dtype not in (jnp.uint32, jnp.int32, jnp.float32):
    raise NotImplementedError(
        f"x.dtype={x.dtype} must be uint32, int32 or float32")
  if not jnp.issubdtype(mask.dtype, jnp.bool):
    raise TypeError(f"mask.dtype={mask.dtype} is not a boolean dtype")
  if x.shape != mask.shape:
    raise ValueError(f"x.shape={x.shape} != mask.shape={mask.shape}")
  return jax_core.ShapedArray(x.shape, jnp.int32), mask


@sc_lowering.register_lowering_rule(scan_count_p)
def _scan_count_lowering_rule(ctx: sc_lowering.LoweringRuleContext, x, mask):
  del ctx  # Unused.
  # Reverse, because the MLIR op returns the mask first.
  return tpu.scan_count(mask, x)[::-1]


def scan_count(
    x: jax.Array, mask: jax.Array | None = None
) -> tuple[jax.Array, jax.Array]:
  """Computes the running duplicate occurrence count of the array.

  Args:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast: mask.astype(jnp.bool_)
2. Keep masks as outputs of comparisons (jnp.bool_ by default)

Example fix

// before
scan_count(x, int_mask)

// after
scan_count(x, int_mask.astype(jnp.bool_))

if not jnp.issubdtype(mask.dtype, jnp.bool_):
    mask = mask.astype(jnp.bool_)

def is_bool_mask(mask) -> bool:
    return jnp.issubdtype(mask.dtype, jnp.bool_)

When it happens

Trigger: scan_count(x, (x > 0).astype(jnp.int32)) or passing a validity bitmap stored as int8.

Common situations: Masks imported as int arrays from HLO dumps, tests, or external data; memory-optimized int masks.

Related errors

• x.dtype={x.dtype} must be uint32, int32 or float32
• x.shape={x.shape} != mask.shape={mask.shape}
• The current TPU does not have SparseCores
• Mesh has {self.num_cores} cores, but the current TPU chip ha
• You can't use two different ScalarSubcoreMeshes.

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