{"record":{"id":"4bea3a1bf779c257","repo":"jax-ml/jax","slug":"name-cannot-accept-args-which-are-reduced-got","errorCode":null,"errorMessage":"{name} cannot accept args which are reduced. Got {a.str_short(True)} and axes={axes}","messagePattern":"(.+?) cannot accept args which are reduced\\. Got (.+?) and axes=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":2758,"sourceCode":"reduced_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:\n    return args\n  in_reduced = [aval.sharding.spec.reduced\n                if isinstance(aval := shaped_abstractify(a), ShapedArray)\n                else frozenset() for a in args]\n  out_reduced = frozenset.union(*in_reduced)\n  out = []\n  for arg, src_reduced in zip(args, in_reduced):\n    aval = shaped_abstractify(arg)\n    if (isinstance(aval, ShapedArray) and aval.ndim == 0 and out_reduced and","sourceCodeStart":2740,"sourceCodeEnd":2776,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L2740-L2776","documentation":"check_unreduced_args' third rule: arguments marked reduced on any of the given axes cannot be consumed by primitives operating on those axes — the data is already reduced there and reusing it as if it varied is invalid.","triggerScenarios":"Using the output of a collective (marked reduced on axis 'x') as input to another primitive over 'x' inside shard_map/spmd without resharding the intermediate.","commonSituations":"Composing psum/all_reduce results into further axis-wise ops; custom spmd code that forwards outputs of collectives directly into more collectives.","solutions":["Reshard the intermediate to clear its 'reduced' marking on those axes","Use the original pre-reduction tensor for follow-up collectives","For custom primitives, declare proper resharding effects"],"exampleFix":"// before\nr = jax.lax.psum(x, 'i')\nr2 = axis_op(r, axis='i')  # r is reduced on 'i'\n\n// after\nr = jax.lax.psum(x, 'i')\nr = reshard(r)\nr2 = axis_op(r, axis='i')","handlingStrategy":"validation","validationCode":"if any(a.mat.reduced & axes for a in args):\n    args = [reshard(a) for a in args]","typeGuard":"def args_clean_of_reduced(args, axes): return not any(a.mat.reduced & set(axes) for a in args)","tryCatchPattern":null,"preventionTips":["Reshard between chained collectives","Feed pre-reduction tensors to follow-up axis ops"],"tags":["jax","sharding","collectives","reduced","spmd"],"backgroundTag":"spmd-axis-misuse","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}