jax-ml/jax · error · ValueError

Expected source shape to be

Error message

Expected source shape to be {expected_src_shape}, but got {src_shape}. {partitioned_desc}

What it means

The abstract eval computes the expected GMEM source shape (from the SMEM destination shape, plus the partitioned dimension's mesh axis size when a collective partition axis is used) and raises ValueError when the actual src shape differs. The message includes which dim is partitioned and its axis size to aid debugging.

Solutions

1. Make the GMEM source shape equal dst.shape, or dst.shape with the partitioned dim multiplied by axis_size if using a partition axis
2. Don't pre-slice the sharded dimension — the collective copy handles partitioning
3. Double-check block_size/grid mappings against the tensor shapes

Example fix

# before
copy_gmem_to_smem(x_sharded, smem, collective_axes=mesh.axes)
# after
copy_gmem_to_smem(x_full, smem, collective_axes=mesh.axes)  # full unsharded dim

expected = list(dst_ref.shape)
if partition_axis is not None:
    expected[partition_axis] *= axis_size
assert list(src_ref.shape) == expected, f'src shape {src_ref.shape} != expected {tuple(expected)}'

When it happens

Trigger: copy_gmem_to_smem where src.shape doesn't match dst.shape (accounting for the partitioned dim being replicated by the collective axis), e.g. a partition axis of size 2 expecting the source to cover the full sharded dimension but receiving a per-shard slice.

Common situations: Sharded/mesh-partitioned kernels where the GMEM ref was already manually sharded but collective_axes expect the full tensor; block-size mismatches between grid and buffers; off-by-one in dimension ordering.

Related errors

• Accumulator shape does not match value shape
• Can only transfer integer bytes
• cluster_names must have the same length as cluster, got
• grid_names must have the same length as grid, got
• Swizzle requires the trailing dimension to be of size , but…

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

Appendix: source

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

    if isinstance(collective_axes, str):
      collective_axes = (collective_axes,)
    axis_size = math.prod(
        jax_core.get_axis_env().axis_size(axis) for axis in collective_axes
    )

  expected_src_shape = tuple(
      d * axis_size if i == partition_axis else d
      for i, d in enumerate(dst_shape)
  )

  if src_shape != expected_src_shape:
    if partition_axis is not None:
      partitioned_desc = (
          f"Dim {partition_axis} is partitioned over an axis of size {axis_size}."
      )
    else:
      partitioned_desc = ""
    raise ValueError(
        f"Expected source shape to be {expected_src_shape}, but got {src_shape}."
        f" {partitioned_desc}"
    )

  return (), {state.ReadEffect(0), state.WriteEffect(1)}


def _copy_gmem_to_smem_pp_eqn(
    eqn: jax_core.JaxprEqn,
    context: jax_core.JaxprPpContext,
    settings: jax_core.JaxprPpSettings,
):
  if eqn.params["has_barrier"]:
    src, dst, barrier, *flat_args = eqn.invars
  else:
    src, dst, *flat_args = eqn.invars
    barrier = None
  src_transforms_treedef = eqn.params["src_transforms_treedef"]

View on GitHub (pinned to 1e1c6a8fc0)