{"record":{"id":"050affc7e75aae09","repo":"jax-ml/jax","slug":"val-dtype-x-dtype-ref-dtype-ref-dtype","errorCode":null,"errorMessage":"val.dtype={x.dtype} != ref.dtype={ref.dtype}","messagePattern":"val\\.dtype=(.+?) != ref\\.dtype=(.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_primitives.py","lineNumber":332,"sourceCode":"scatter_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)\ndef _scatter_lowering_rule(\n    ctx: sc_lowering.LoweringRuleContext, *flat_args, tree, add\n):\n  ref, transforms, indices, x, mask = jax.tree.unflatten(tree, flat_args)","sourceCodeStart":314,"sourceCodeEnd":350,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_primitives.py#L314-L350","documentation":"The value scattered into the ref must have exactly the same dtype as the ref itself. Unlike standard JAX ops, no implicit dtype promotion happens in SC scatter.","triggerScenarios":"x is e.g. float32 while the ref was declared int32 (or vice versa), or bf16 value into f32 ref.","commonSituations":"Mixing computation dtype (often bf16/f32) with a differently-typed output buffer; forgetting that Pallas requires exact dtype agreement.","solutions":["Cast x to ref.dtype before the scatter (x.astype(ref_dtype))","Or declare the ref with the same dtype as your computed values"],"exampleFix":"// before\nsc_primitives.store_scatter(ref, indices, x)  # x f32, ref int32\n\n// after\nsc_primitives.store_scatter(ref, indices, x.astype(ref.dtype))","handlingStrategy":"type-guard","validationCode":"if x.dtype != ref.dtype:\n    x = x.astype(ref.dtype)","typeGuard":"def dtypes_match(x, ref) -> bool:\n    return x.dtype == getattr(getattr(ref, 'aval', ref), 'dtype', None)","tryCatchPattern":null,"preventionTips":["Always cast x to ref.dtype right before scatter","Avoid mixing compute dtype and buffer dtype in SC kernels","Add asserts in kernel tests covering dtype agreement"],"tags":["jax","pallas","sparsecore","scatter","dtype-mismatch"],"backgroundTag":"dtype-mismatch-assignment","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}