{"record":{"id":"309d1ba95eecaf6f","repo":"jax-ml/jax","slug":"name-cannot-accept-args-which-are-unreduced-got","errorCode":null,"errorMessage":"{name} cannot accept args which are unreduced. Got {a.str_short(True)} and axes={axes}","messagePattern":"(.+?) cannot accept args which are unreduced\\. Got (.+?) and axes=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":2749,"sourceCode":"    return x\n  cur_mesh = mesh_lib.get_abstract_mesh()\n  if not config._check_vma.value and all(a in cur_mesh.manual_axes for a in axes):\n    return x\n  new_axes = axes if cur_mesh.empty else order_wrt_mesh(cur_mesh, axes)\n  assert set(new_axes) == set(axes)\n  del axes\n  return tree_map(lambda leaf: reduced_vary_cast_p.bind(leaf, axes=new_axes), x)\n\nreduced_vary_cast_p = Primitive('reduced_vary_cast_p')\n\n#######################################################################\n\ndef check_unreduced_args(args, axes, name, kind=UnreducedKind.sum):\n  axes = axes if isinstance(axes, (tuple, list)) else (axes,)\n  axes = set(axes)\n  for a in args:\n    if a.mat.unreduced & axes:\n      raise ValueError(\n          f\"{name} cannot accept args which are unreduced. Got\"\n          f\" {a.str_short(True)} and axes={axes}\")\n    if a.mat.unreduced and a.mat.unreduced_kind is not kind:\n      raise ValueError(\n          f\"{name} cannot accept args with\"\n          f\" unreduced_kind={a.mat.unreduced_kind}. Expected\"\n          f\" unreduced_kind={kind}\")\n    if a.mat.reduced & axes:\n      raise ValueError(\n          f\"{name} cannot accept args which are reduced. Got\"\n          f\" {a.str_short(True)} and axes={axes}\")\n\ndef insert_reduced_reshard(args):\n  cur_mesh = mesh_lib.get_abstract_mesh()\n  if not cur_mesh.are_all_axes_explicit:\n    return args\n  # TODO(yashkatariya): Handle >2 args too\n  if len(args) != 2:","sourceCodeStart":2731,"sourceCodeEnd":2767,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L2731-L2767","documentation":"check_unreduced_args guards collective/axis-consuming primitives: arguments marked unreduced on any of the given mesh axes cannot be consumed there, because the value is only valid un-reduced on those axes. E.g. a psum-intermediate consumed again along its reduction axis is rejected.","triggerScenarios":"Inside shard_map/spmd code, feeding an array whose mat.unreduced intersects the axes a primitive operates over — e.g. using an unreduced psum output in another collective over the same axis without resharding.","commonSituations":"Chaining collectives (psum output reused in another psum/all_gather) on the same mesh axis; writing custom spmd primitives without inserting resharding.","solutions":["Reshard/convert the value so it is no longer unreduced on those axes (e.g. insert_reduced_reshard or an explicit reshard) before consuming it","Restructure the computation to consume the pre-reduction value or use the reduced result instead","For custom primitives, mark the operand axes properly so the framework inserts resharding"],"exampleFix":"// before\ny = jax.lax.psum(x, 'i')   # x unreduced on 'i'\nz = jax.lax.psum(y, 'i')   # consuming unreduced on same axis\n\n// after\ny = jax.lax.psum(x, 'i')\ny = reshard(y)             # clear unreduced state on 'i'\nz = some_local_op(y)","handlingStrategy":"validation","validationCode":"if any(a.mat.unreduced & axes for a in args):\n    args = [reshard(a) for a in args]","typeGuard":"def args_clean_of_unreduced(args, axes): return not any(a.mat.unreduced & set(axes) for a in args)","tryCatchPattern":null,"preventionTips":["Reshard collective outputs before reuse","Keep psum results flowing to replicated ops only"],"tags":["jax","sharding","collectives","unreduced","spmd"],"backgroundTag":"spmd-axis-misuse","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}