{"record":{"id":"eb982922296e7b65","repo":"jax-ml/jax","slug":"expected-source-shape-to-be-expected-src-shape","errorCode":null,"errorMessage":"Expected source shape to be {expected_src_shape}, but got {src_shape}. {partitioned_desc}","messagePattern":"Expected source shape to be (.+?), but got (.+?)\\. (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":813,"sourceCode":"    if isinstance(collective_axes, str):\n      collective_axes = (collective_axes,)\n    axis_size = math.prod(\n        jax_core.get_axis_env().axis_size(axis) for axis in collective_axes\n    )\n\n  expected_src_shape = tuple(\n      d * axis_size if i == partition_axis else d\n      for i, d in enumerate(dst_shape)\n  )\n\n  if src_shape != expected_src_shape:\n    if partition_axis is not None:\n      partitioned_desc = (\n          f\"Dim {partition_axis} is partitioned over an axis of size {axis_size}.\"\n      )\n    else:\n      partitioned_desc = \"\"\n    raise ValueError(\n        f\"Expected source shape to be {expected_src_shape}, but got {src_shape}.\"\n        f\" {partitioned_desc}\"\n    )\n\n  return (), {state.ReadEffect(0), state.WriteEffect(1)}\n\n\ndef _copy_gmem_to_smem_pp_eqn(\n    eqn: jax_core.JaxprEqn,\n    context: jax_core.JaxprPpContext,\n    settings: jax_core.JaxprPpSettings,\n):\n  if eqn.params[\"has_barrier\"]:\n    src, dst, barrier, *flat_args = eqn.invars\n  else:\n    src, dst, *flat_args = eqn.invars\n    barrier = None\n  src_transforms_treedef = eqn.params[\"src_transforms_treedef\"]","sourceCodeStart":795,"sourceCodeEnd":831,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L795-L831","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Make the GMEM source shape equal dst.shape, or dst.shape with the partitioned dim multiplied by axis_size if using a partition axis","Don't pre-slice the sharded dimension — the collective copy handles partitioning","Double-check block_size/grid mappings against the tensor shapes"],"exampleFix":"# before\ncopy_gmem_to_smem(x_sharded, smem, collective_axes=mesh.axes)\n# after\ncopy_gmem_to_smem(x_full, smem, collective_axes=mesh.axes)  # full unsharded dim","handlingStrategy":"validation","validationCode":"expected = list(dst_ref.shape)\nif partition_axis is not None:\n    expected[partition_axis] *= axis_size\nassert list(src_ref.shape) == expected, f'src shape {src_ref.shape} != expected {tuple(expected)}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Don't pre-shard the GMEM dim covered by the partition axis","Keep block_size, grid, and tensor shapes derived from one source of truth"],"tags":["jax","pallas","mosaic-gpu","shape-mismatch","sharding","gmem-copy"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}