{"record":{"id":"16f83fe59a0d5952","repo":"jax-ml/jax","slug":"unreduced-psum-scatter-is-a-unreduced-varying-c","errorCode":null,"errorMessage":"unreduced_psum_scatter is a Unreduced -> Varying collective. This means that the {axis_name=} passed to `unreduced_psum_scatter` must be present in jax.typeof(x).mat.unreduced={x_aval.mat.unreduced}","messagePattern":"unreduced_psum_scatter is a Unreduced -> Varying collective\\. This means that the (.+?) passed to `unreduced_psum_scatter` must be present in jax\\.typeof\\(x\\)\\.mat\\.unreduced=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2677,"sourceCode":"  axis_size = _axis_size(axis_name, None)\n  def bind(leaf):\n    return unreduced_reduce_scatter_p.bind(\n        leaf, axis_name=axis_name, scatter_dimension=scatter_dimension,\n        axis_size=axis_size, tiled=tiled)\n  return tree_util.tree_map(bind, x)\n\nunreduced_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","sourceCodeStart":2659,"sourceCodeEnd":2695,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2659-L2695","documentation":"The `axis_name` passed to `unreduced_psum_scatter` must be present in `jax.typeof(x).mat.unreduced`; if the input is unreduced on other axes but not the requested one, JAX raises this error. The input's unreduced state and the scatter axis must agree. It usually means a mismatch between mesh axis names or a missing unreduced annotation.","triggerScenarios":"x is unreduced along 'i' but calling with `axis_name='j'`; or passing axis names that do not intersect `x.mat.unreduced`.","commonSituations":"Copy-pasting collective calls between model shards with different mesh axis names; typos in axis_name; renaming axes in the mesh without updating collective calls.","solutions":["Print `jax.typeof(x).mat.unreduced` and pass one of those axes as axis_name","Add the missing unreduced producer for the intended axis","Verify axis names against the mesh/annotation definitions for typos"],"exampleFix":"// before\ny = lax.unreduced_psum_scatter(x, axis_name='j', ...)\n// after\ny = lax.unreduced_psum_scatter(x, axis_name='i', ...)  # 'i' in x.mat.unreduced","handlingStrategy":"validation","validationCode":"unreduced = jax.typeof(x).mat.unreduced\nassert set(unreduced) & set(axis_name), f'{axis_name} not in {unreduced}'\ny = lax.unreduced_psum_scatter(x, axis_name=axis_name, ...)","typeGuard":"def axis_is_unreduced(x, axis_name) -> bool:\n    return bool(set(jax.typeof(x).mat.unreduced) & set(axis_name))","tryCatchPattern":"try:\n    y = lax.unreduced_psum_scatter(x, axis_name=axis_name, ...)\nexcept ValueError as e:\n    if 'Unreduced -> Varying' in str(e):\n        axis_name = tuple(jax.typeof(x).mat.unreduced)[:1]\n        y = lax.unreduced_psum_scatter(x, axis_name=axis_name, ...)\n    else:\n        raise","preventionTips":["Use shared constants for mesh axis names","Assert axis membership in mat.unreduced before scattering"],"tags":["jax","collectives","axis-name","mesh","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"}