{"record":{"id":"5359c6a77f1086aa","repo":"jax-ml/jax","slug":"x-shape-x-shape-mask-shape-mask-shape","errorCode":null,"errorMessage":"x.shape={x.shape} != mask.shape={mask.shape}","messagePattern":"x\\.shape=(.+?) != mask\\.shape=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_primitives.py","lineNumber":531,"sourceCode":"  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]:\n  \"\"\"Computes the running duplicate occurrence count of the array.\n\n  Args:\n    x: An array of integers or floats.\n    mask: An optional array of booleans, which specifies which elements ``x``","sourceCodeStart":513,"sourceCodeEnd":549,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_primitives.py#L513-L549","documentation":"scan_count requires mask.shape == x.shape exactly; broadcasting is not applied.","triggerScenarios":"Passing mask of shape (N,) with x of shape (N, M), or a scalar mask with any x.","commonSituations":"Assuming mask broadcasts over the vector dimension; reusing a batch-level mask for token-level data.","solutions":["Broadcast the mask to x.shape explicitly: jnp.broadcast_to(mask, x.shape)","Recompute the mask at matching shape inside the kernel"],"exampleFix":"// before\nscan_count(x, mask)  # mask (N,), x (N,M)\n\n// after\nscan_count(x, jnp.broadcast_to(mask, x.shape))","handlingStrategy":"validation","validationCode":"if mask.shape != x.shape:\n    mask = jnp.broadcast_to(mask, x.shape)","typeGuard":"def shapes_match(x, mask) -> bool:\n    return tuple(x.shape) == tuple(mask.shape)","tryCatchPattern":null,"preventionTips":["Broadcast masks explicitly; never rely on implicit rules","Derive masks from the tensor they mask","Test with non-trivial shapes"],"tags":["jax","pallas","sparsecore","scan-count","shape-mismatch"],"backgroundTag":"operand-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}