{"record":{"id":"d0d62c11b07948b1","repo":"jax-ml/jax","slug":"name-cannot-accept-args-with-unreduced-kind-mat","errorCode":null,"errorMessage":"{name} cannot accept args with unreduced_kind={mat.unreduced_kind}. Expected unreduced_kind={kind}","messagePattern":"(.+?) cannot accept args with unreduced_kind=(.+?)\\. Expected unreduced_kind=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":64,"sourceCode":"from jax._src.lax import lax\nfrom jax._src.lax import slicing\nfrom jax._src.lib.mlir import ir\nfrom jax._src.lib.mlir.dialects import hlo\nfrom jax._src.typing import Array\nfrom jax._src.util import (canonicalize_axis, moveaxis, safe_map, safe_zip,\n                           unzip2)\nfrom jax._src.lib.mlir.dialects import func as func_dialect\nimport numpy as np\n\nunsafe_map, map = map, safe_map\nunsafe_zip, zip = zip, safe_zip\n\n\n### parallel traceables\n\ndef check_unreduced_kind(name, mat, kind):\n  if mat.unreduced_kind is not kind:\n    raise ValueError(\n        f\"{name} cannot accept args with unreduced_kind={mat.unreduced_kind}.\"\n        f\" Expected unreduced_kind={kind}\")\n\ndef psum(x, axis_name, *, axis_index_groups=None):\n  \"\"\"Compute an all-reduce sum on ``x`` over the pmapped axis ``axis_name``.\n\n  If ``x`` is a pytree then the result is equivalent to mapping this function to\n  each leaf in the tree.\n\n  Inputs of boolean dtype are converted to integers before the reduction.\n\n  Args:\n    x: array(s) with a mapped axis named ``axis_name``.\n    axis_name: hashable Python object used to name a pmapped axis (see the\n      :func:`jax.pmap` documentation for more details).\n    axis_index_groups: optional list of lists containing axis indices (e.g. for\n      an axis of size 4, [[0, 1], [2, 3]] would perform psums over the first\n      two and last two replicas). Groups must cover all axis indices exactly","sourceCodeStart":46,"sourceCodeEnd":82,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L46-L82","documentation":"jax.lax parallel collectives operate on arrays that may be 'unreduced' or 'reduced' with respect to a mapped axis (used by shard_map / SPMD typed collectives). check_unreduced_kind guards each primitive so it only receives the kind of array it can semantically handle — e.g. all_gather's reduced-effectful path expects unreduced inputs. A mismatch means the collective was applied to an array whose reduction status under that axis does not fit the operation.","triggerScenarios":"Mixing typed SPMD collectives, e.g. calling the reduced all_gather / psum / pmax path on an array already marked unreduced (or vice versa) for axis_name — typically inside shard_map pipelines composing psum, all_gather, and reduce_scatter in one function.","commonSituations":"Refactoring shard_map code where a previous collective changed the array's unreduced_kind; upgrading JAX versions where typed-collective rules tightened; hand-writing compose pipelines of psum -> all_gather -> reduce_scatter.","solutions":["Re-check the order of collectives in your shard_map body; ensure the input to each collective has the expected reduction state (e.g. don't feed all_gather-reduced output back into the unreduced path)","Use plain lax.psum / all_gather traceables instead of lower-level typed primitives so JAX tracks state for you","Upgrade/downgrade to a JAX version whose shard_map semantics match the code (these rules changed across releases)"],"exampleFix":"// before\ny = all_gather_reduced(psum_out, axis_name)  # psum output is reduced -> mismatch\n\n// after\ny = jax.lax.all_gather(x, axis_name)  # operate on the original unreduced x","handlingStrategy":"validation","validationCode":"# prefer public traceables (jax.lax.psum etc.) so unreduced_kind is tracked automatically","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid low-level typed-collective primitives directly","Keep collective order simple: one reduce per axis per shard_map body"],"tags":["jax","shard-map","spmd","collectives","internal-api"],"backgroundTag":"invalid-collective-usage","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}