jax-ml/jax · error · TypeError

val.dtype={x.dtype} != ref.dtype={ref.dtype}

Error message

val.dtype={x.dtype} != ref.dtype={ref.dtype}

What it means

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.

Source

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

scatter_p = jax_core.Primitive("scatter")
scatter_p.is_effectful = lambda params: True
scatter_p.multiple_results = True


@scatter_p.def_effectful_abstract_eval
def _scatter_abstract_eval(*flat_args, tree, add):
  ref, transforms, indices, x, mask = jax.tree.unflatten(tree, flat_args)
  if transforms:
    ref = state_types.TransformedRef(ref, transforms)
  if ref.dtype not in (jnp.int32, jnp.float32):
    raise TypeError(f"ref.dtype={ref.dtype} must be int32 or float32")
  expected_shape = _indexed_shape(ref, indices)
  if x.shape != expected_shape:
    raise ValueError(
        f"{x.shape=} does not match expected shape {expected_shape}"
    )
  if x.dtype != ref.dtype:
    raise TypeError(f"val.dtype={x.dtype} != ref.dtype={ref.dtype}")
  if mask is not None:
    if mask.shape != expected_shape:
      raise ValueError(
          f"{mask.shape=} does not match expected shape {expected_shape}"
      )
    if mask.dtype != jnp.bool:
      raise TypeError(f"Mask must be a boolean array, got {mask.dtype}")
  effects: set[jax_core.Effect] = {state_types.WriteEffect(0)}
  if add:
    effects.add(state_types.ReadEffect(0))
  return (), effects


@sc_lowering.register_lowering_rule(scatter_p)
def _scatter_lowering_rule(
    ctx: sc_lowering.LoweringRuleContext, *flat_args, tree, add
):
  ref, transforms, indices, x, mask = jax.tree.unflatten(tree, flat_args)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast x to ref.dtype before the scatter (x.astype(ref_dtype))
2. Or declare the ref with the same dtype as your computed values

Example fix

// before
sc_primitives.store_scatter(ref, indices, x)  # x f32, ref int32

// after
sc_primitives.store_scatter(ref, indices, x.astype(ref.dtype))

if x.dtype != ref.dtype:
    x = x.astype(ref.dtype)

def dtypes_match(x, ref) -> bool:
    return x.dtype == getattr(getattr(ref, 'aval', ref), 'dtype', None)

When it happens

Trigger: x is e.g. float32 while the ref was declared int32 (or vice versa), or bf16 value into f32 ref.

Common situations: Mixing computation dtype (often bf16/f32) with a differently-typed output buffer; forgetting that Pallas requires exact dtype agreement.

Related errors

• {x.shape=} does not match expected shape {expected_shape}
• {mask.shape=} does not match expected shape {expected_shape}
• Scatter only supports VectorSubcoreMesh, got {type(ref_aval.
• Scatter only supports storing to VMEM, got {memory_space}
• Indices must not be empty

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