jax-ml/jax · error · TypeError

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

Error message

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

What it means

Abstract eval of async_store_smem compares the shape of the value being stored (src.shape) with the shape of the target SMEM reference after applying its transforms (TransformedRef). If they differ, a TypeError is raised because the async hardware store copies raw bits and cannot reshape.

Source

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

    barrier_transforms_treedef,
    **_,
):
  del cluster_idx  # Unused.
  _check_ref(ref, "ref", gpu_core.SMEM)
  _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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make src.shape exactly equal transformed_ref.shape (adjust block_size or reshape the value first)
2. Apply matching index/slice transforms to the SMEM ref so its transformed shape equals the value's shape
3. Store via a regular store (smem[...] = value) if reshaping semantics are needed

Example fix

# before
async_store_smem(smem, x, barrier)  # x.shape=(128,), smem.shape=(16,8)
# after
async_store_smem(smem, x.reshape(smem.shape), barrier)

assert src.shape == tuple(t_ref.shape), (
  f'async_store_smem shape mismatch: {src.shape} vs {t_ref.shape}')

When it happens

Trigger: Calling async_store_smem(smem_ref, value, ...) where value.shape != smem_ref.shape after slicing/tiling transforms; e.g. storing a (128,) vector into a (16,8) tiled ref, or omitting a slice that reshapes the target.

Common situations: Off-by-one block sizes between compute and buffers; forgetting to apply the same slice to the ref as used for the source; changing block shapes without resizing SMEM buffers.

Related errors

• Swizzle {self.swizzle} requires the trailing dimension to be
• grid_names must have the same length as grid, got {self}.
• cluster_names must have the same length as cluster, got {sel
• {tiling=} and {grid=} must have same length.
• The stored value has dtype {src.dtype}, but the target refer

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