{"record":{"id":"1e9f777d8f7617c6","repo":"jax-ml/jax","slug":"x-dtype-x-dtype-must-be-uint32-int32-or-float32","errorCode":null,"errorMessage":"x.dtype={x.dtype} must be uint32, int32 or float32","messagePattern":"x\\.dtype=(.+?) must be uint32, int32 or float32","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_primitives.py","lineNumber":526,"sourceCode":"\n\ndef subcore_barrier():\n  \"\"\"Blocks until all subcores on the same core reach this instruction.\n\n  The barrier must be used with\n  :class:`jax.experimental.pallas.tpu_sc.VectorSubcoreMesh`.\n  \"\"\"\n  barrier_p.bind()\n\n\nscan_count_p = jax_core.Primitive(\"scan_count\")\nscan_count_p.multiple_results = True\n\n\n@scan_count_p.def_abstract_eval\ndef _scan_count_abstract_eval(x, mask):\n  if x.dtype not in (jnp.uint32, jnp.int32, jnp.float32):\n    raise NotImplementedError(\n        f\"x.dtype={x.dtype} must be uint32, int32 or float32\")\n  if not jnp.issubdtype(mask.dtype, jnp.bool):\n    raise TypeError(f\"mask.dtype={mask.dtype} is not a boolean dtype\")\n  if x.shape != mask.shape:\n    raise ValueError(f\"x.shape={x.shape} != mask.shape={mask.shape}\")\n  return jax_core.ShapedArray(x.shape, jnp.int32), mask\n\n\n@sc_lowering.register_lowering_rule(scan_count_p)\ndef _scan_count_lowering_rule(ctx: sc_lowering.LoweringRuleContext, x, mask):\n  del ctx  # Unused.\n  # Reverse, because the MLIR op returns the mask first.\n  return tpu.scan_count(mask, x)[::-1]\n\n\ndef scan_count(\n    x: jax.Array, mask: jax.Array | None = None\n) -> tuple[jax.Array, jax.Array]:","sourceCodeStart":508,"sourceCodeEnd":544,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_primitives.py#L508-L544","documentation":"The SparseCore scan_count primitive only supports uint32, int32, or float32 input; other dtypes raise NotImplementedError at abstract eval.","triggerScenarios":"Calling scan_count(x, mask) with x of dtype bfloat16, int64, float16, etc.","commonSituations":"Feeding bf16 activations (common in attention/embedding pipelines) into scan_count; using default int64 indices as the counted array.","solutions":["Cast x to uint32/int32/float32 before scan_count","If x is a mask-like count, use .astype(jnp.int32)","Check upstream dtype of the tensor feeding scan_count in the kernel"],"exampleFix":"// before\ncnt, m = scan_count(x_bf16, mask)\n\n// after\ncnt, m = scan_count(x_bf16.astype(jnp.float32), mask)","handlingStrategy":"validation","validationCode":"if x.dtype not in (jnp.uint32, jnp.int32, jnp.float32):\n    x = x.astype(jnp.float32)","typeGuard":"def scan_count_dtype_ok(x) -> bool:\n    return x.dtype in (jnp.dtype('uint32'), jnp.dtype('int32'), jnp.dtype('float32'))","tryCatchPattern":null,"preventionTips":["Cast bf16/int64 inputs to f32/i32 at SC kernel entry","Wrap scan_count in a helper that normalizes dtypes","Watch for int64 defaults from arange/constants"],"tags":["jax","pallas","sparsecore","scan-count","dtype"],"backgroundTag":"unsupported-dtype-on-accelerator-op","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}