{"record":{"id":"8bec0fc3cabb5cd0","repo":"jax-ml/jax","slug":"x-shape-does-not-match-expected-shape-expected","errorCode":null,"errorMessage":"{x.shape=} does not match expected shape {expected_shape}","messagePattern":"(.+?) does not match expected shape (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_primitives.py","lineNumber":328,"sourceCode":"  flat_args, tree = jax.tree.flatten((ref, transforms, indices, mask))\n  return gather_p.bind(*flat_args, tree=tree)\n\n\nscatter_p = jax_core.Primitive(\"scatter\")\nscatter_p.is_effectful = lambda params: True\nscatter_p.multiple_results = True\n\n\n@scatter_p.def_effectful_abstract_eval\ndef _scatter_abstract_eval(*flat_args, tree, add):\n  ref, transforms, indices, x, mask = jax.tree.unflatten(tree, flat_args)\n  if transforms:\n    ref = state_types.TransformedRef(ref, transforms)\n  if ref.dtype not in (jnp.int32, jnp.float32):\n    raise TypeError(f\"ref.dtype={ref.dtype} must be int32 or float32\")\n  expected_shape = _indexed_shape(ref, indices)\n  if x.shape != expected_shape:\n    raise ValueError(\n        f\"{x.shape=} does not match expected shape {expected_shape}\"\n    )\n  if x.dtype != ref.dtype:\n    raise TypeError(f\"val.dtype={x.dtype} != ref.dtype={ref.dtype}\")\n  if mask is not None:\n    if mask.shape != expected_shape:\n      raise ValueError(\n          f\"{mask.shape=} does not match expected shape {expected_shape}\"\n      )\n    if mask.dtype != jnp.bool:\n      raise TypeError(f\"Mask must be a boolean array, got {mask.dtype}\")\n  effects: set[jax_core.Effect] = {state_types.WriteEffect(0)}\n  if add:\n    effects.add(state_types.ReadEffect(0))\n  return (), effects\n\n\n@sc_lowering.register_lowering_rule(scatter_p)","sourceCodeStart":310,"sourceCodeEnd":346,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_primitives.py#L310-L346","documentation":"Scatter's abstract eval computes the expected shape of the value from the ref shape and the indices (via _indexed_shape) and the provided value array x does not match. The value must exactly fill the region selected by the indices.","triggerScenarios":"Passing an x whose shape differs from _indexed_shape(ref, indices), e.g. indices of shape (N,) selecting rows of size M but x shaped (N, K) with K != M, or a mismatched leading dimension.","commonSituations":"Off-by-one in block sizes; broadcasting assumptions from NumPy that Pallas does not make; changing indices semantics (per-element vs per-slice) when porting code.","solutions":["Print/inspect _indexed_shape(ref, indices) (shape arithmetic: ref shape minus indexed dims) and reshape x to match","Fix the index shapes so the selected region matches x.shape","Explicitly broadcast/reshape x before calling store_scatter/addupdate_scatter"],"exampleFix":"// before\nsc_primitives.store_scatter(ref, idx, x)  # x.shape=(8, 32), expected=(8, 16)\n\n// after\nsc_primitives.store_scatter(ref, idx, x.reshape(8, 16))  # or fix idx/block shapes","handlingStrategy":"validation","validationCode":"expected = _indexed_shape(ref, indices)  # or reimplement: ref.shape minus indexed dims\nassert x.shape == expected, (x.shape, expected)","typeGuard":"def scatter_shapes_ok(ref, indices, x) -> bool:\n    indexed = sum(i.shape[-1] if hasattr(i, 'shape') else 1 for i in indices)\n    return tuple(x.shape) == tuple(ref.shape[indexed:])","tryCatchPattern":null,"preventionTips":["Write shape preconditions next to each scatter call","Print expected vs actual shapes in kernel debug runs","Keep index count consistent with x's leading dims"],"tags":["jax","pallas","sparsecore","scatter","shape-mismatch"],"backgroundTag":"shape-mismatch-scatter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}