jax-ml/jax · error · TypeError

The stored value has dtype {src.dtype}, but the target refer

Error message

The stored value has dtype {src.dtype}, but the target reference has dtype {transformed_ref.dtype}

What it means

async_store_smem's abstract eval checks that the dtype of the stored value matches the dtype of the transformed target SMEM reference. Mismatch raises TypeError because the async copy is bit-exact with no implicit casting.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:580

  _check_ref(barrier, "barrier", gpu_core.SMEM)
  flat_ref_transforms_avals, flat_barrier_transforms_avals = util.split_list(
      flat_transforms_avals,
      [ref_transforms_treedef.num_leaves],
  )
  ref_transform_avals = ref_transforms_treedef.unflatten(
      flat_ref_transforms_avals
  )
  barrier_transform_avals = barrier_transforms_treedef.unflatten(
      flat_barrier_transforms_avals
  )
  transformed_ref = pallas_core.TransformedRef(ref, ref_transform_avals)
  if src.shape != transformed_ref.shape:
    raise TypeError(
        f"The stored value has shape {src.shape}, but the target reference has"
        f" shape {transformed_ref.shape}"
    )
  if src.dtype != transformed_ref.dtype:
    raise TypeError(
        f"The stored value has dtype {src.dtype}, but the target reference has"
        f" dtype {transformed_ref.dtype}"
    )
  transformed_barrier = pallas_core.TransformedRef(barrier, barrier_transform_avals)
  if transformed_barrier.size != 1:
    raise TypeError(
        "Expected a single barrier, got a barrier reference with shape"
        f" {transformed_barrier.shape}"
    )

  effs = {gpu_core._memory_effect, state.WriteEffect(1)}
  return (), effs


@lowering.register_lowering_rule(async_store_smem_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(async_store_smem_p, mgpu.LoweringSemantics.Warpgroup)
def _async_store_smem_lowering(
    ctx: lowering.LoweringRuleContext,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate the SMEM block with the same dtype as the value (e.g. pl.SMEM((128,), x.dtype))
2. Explicitly cast the value before storing: x.astype(smem.dtype)
3. If a wider accumulation type is intended, use two stores or a regular store with a cast

Example fix

# before
smem = pl.SMEM((128,), jnp.float32)
async_store_smem(smem, x_bf16, barrier)
# after
async_store_smem(smem, x_bf16.astype(jnp.float32), barrier)

assert src.dtype == smem_ref.dtype, (
  f'async_store_smem dtype mismatch: {src.dtype} vs {smem_ref.dtype}')

When it happens

Trigger: Calling async_store_smem where src.dtype differs from the SMEM buffer's dtype, e.g. storing a float32 tensor into a bfloat16 SMEM block, or int32 into int8.

Common situations: Allocating SMEM scratch without an explicit dtype so it defaults differently from the computed value; mixing precision regimes (bf16 compute, fp32 buffers) in fused kernels; refactoring a kernel to half precision but keeping old scratch dtypes.

Related errors

• The stored value has shape {src.shape}, but the target refer
• Expected a single barrier, got a barrier reference with shap
• Scalars are not supported in async_store_smem
• Unexpected unhandled transforms: {remaining_ref_transforms}
• async_store_smem requires a tiled and swizzled ref

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