{"record":{"id":"d492a34230b2c7a7","repo":"jax-ml/jax","slug":"unreduced-psum-scatter-s-input-cannot-be-varying-a","errorCode":null,"errorMessage":"unreduced_psum_scatter's input cannot be varying across the axis_name provided. Got x={x_aval.str_short(True)} and {axis_name=}","messagePattern":"unreduced_psum_scatter's input cannot be varying across the axis_name provided\\. Got x=(.+?) and (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2684,"sourceCode":"unreduced_reduce_scatter_p = core.Primitive('unreduced_reduce_scatter')\n\ndef _unreduced_reduce_scatter_effectful_abstract_eval(\n    x_aval, *, axis_name, scatter_dimension, axis_size, tiled\n):\n  _check_axis_names(axis_name, 'reduce_scatter')\n  if not x_aval.mat.unreduced:\n    raise ValueError('unreduced_psum_scatter only accepts inputs that are'\n                     f' unreduced. Got {x_aval.str_short(True)}')\n  # If intersection between x.unreduced & axis_name is empty, error\n  if not (x_aval.mat.unreduced & frozenset(axis_name)):\n    raise ValueError(\n        \"unreduced_psum_scatter is a Unreduced -> Varying collective. This\"\n        f\" means that the {axis_name=} passed to `unreduced_psum_scatter` must\"\n        \" be present in\"\n        f\" jax.typeof(x).mat.unreduced={x_aval.mat.unreduced}\"\n    )\n  if x_aval.mat.varying & set(axis_name):\n    raise ValueError(\n        \"unreduced_psum_scatter's input cannot be varying across the axis_name\"\n        f\" provided. Got x={x_aval.str_short(True)} and {axis_name=}\")\n\n  new_shape = list(x_aval.shape)\n  scatter_dim_input_size = x_aval.shape[scatter_dimension]\n  if tiled:\n    if scatter_dim_input_size % axis_size != 0:\n      raise ValueError(f\"tiled reduce_scatter operand scatter dimension size \"\n                       f\"{scatter_dim_input_size} must be divisible by \"\n                       f\"shard_count {axis_size}\")\n    new_shape[scatter_dimension] = scatter_dim_input_size // axis_size\n  else:\n    if scatter_dim_input_size != axis_size:\n      raise ValueError(f\"reduce_scatter operand scatter dimension size \"\n                       f\"{scatter_dim_input_size} must match shard count \"\n                       f\"{axis_size}\")\n    del new_shape[scatter_dimension]\n","sourceCodeStart":2666,"sourceCodeEnd":2702,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2666-L2702","documentation":"`unreduced_psum_scatter`'s input must not already be varying along the requested `axis_name` (i.e. `x.mat.varying` must not intersect it). A value cannot be simultaneously varying and about to become varying via scatter on the same axis. The abstract eval rejects this inconsistent state.","triggerScenarios":"Input previously marked with pvary along 'i' then passed to `unreduced_psum_scatter(..., axis_name='i')`.","commonSituations":"Composing pvary with reduce-scatter pipelines where the same axis ends up in both varying and unreduced tracks; porting older sharding code that applied explicit sharding annotations before collective calls.","solutions":["Remove the earlier pvary on that axis, or scatter along a different axis not in `x.mat.varying`","Restructure so the value is unreduced (not varying) on the scatter axis before the call","Inspect `jax.typeof(x).mat` to see the full varying/unreduced state and fix the producer chain"],"exampleFix":"// before\nx = lax.pvary(x, 'i')\ny = lax.unreduced_psum_scatter(x, axis_name='i', ...)\n// after\ny = lax.unreduced_psum_scatter(unreduced_x, axis_name='i', ...)  # x not varying on 'i'","handlingStrategy":"validation","validationCode":"assert not (jax.typeof(x).mat.varying & set(axis_name)), (\n    f'{axis_name} already varying in {jax.typeof(x).mat.varying}')\ny = lax.unreduced_psum_scatter(x, axis_name=axis_name, ...)","typeGuard":"def not_varying_on(x, axis_name) -> bool:\n    return not (jax.typeof(x).mat.varying & set(axis_name))","tryCatchPattern":"try:\n    y = lax.unreduced_psum_scatter(x, axis_name=axis_name, ...)\nexcept ValueError as e:\n    if 'cannot be varying across' in str(e):\n        raise RuntimeError(f'state conflict on {axis_name}; fix producer chain') from e\n    raise","preventionTips":["Document which pipeline stage owns each axis's state (varying/unreduced)","Avoid mixing pvary annotations with unreduced-sum pipelines on the same axis"],"tags":["jax","collectives","mesh","varying","reduce-scatter"],"backgroundTag":"jax-mat-axis-state-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}