jax-ml/jax · error · NotImplementedError

x.dtype={x.dtype} must be uint32, int32 or float32

Error message

x.dtype={x.dtype} must be uint32, int32 or float32

What it means

The SparseCore scan_count primitive only supports uint32, int32, or float32 input; other dtypes raise NotImplementedError at abstract eval.

Source

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

def subcore_barrier():
  """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]:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast x to uint32/int32/float32 before scan_count
2. If x is a mask-like count, use .astype(jnp.int32)
3. Check upstream dtype of the tensor feeding scan_count in the kernel

Example fix

// before
cnt, m = scan_count(x_bf16, mask)

// after
cnt, m = scan_count(x_bf16.astype(jnp.float32), mask)

if x.dtype not in (jnp.uint32, jnp.int32, jnp.float32):
    x = x.astype(jnp.float32)

def scan_count_dtype_ok(x) -> bool:
    return x.dtype in (jnp.dtype('uint32'), jnp.dtype('int32'), jnp.dtype('float32'))

When it happens

Trigger: Calling scan_count(x, mask) with x of dtype bfloat16, int64, float16, etc.

Common situations: Feeding bf16 activations (common in attention/embedding pipelines) into scan_count; using default int64 indices as the counted array.

Related errors

• Unsupported dtype: {dtype}
• mask.dtype={mask.dtype} is not a boolean dtype
• x.shape={x.shape} != mask.shape={mask.shape}
• Index map function {debug_info.func_src_info} for {origin} m
• bitcast_convert_type with different bitwidths not supported

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