{"record":{"id":"a327e04225bb9fa8","repo":"jax-ml/jax","slug":"varying-and-reduced-cannot-have-common-mesh-axes","errorCode":null,"errorMessage":"varying and reduced cannot have common mesh axes. Got varying={varying} and reduced={reduced}","messagePattern":"varying and reduced cannot have common mesh axes\\. Got varying=(.+?) and reduced=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":2349,"sourceCode":"    reduced = frozenset(r for r in mat.reduced if mesh.shape[r] != 1)\n    u_kind = mat.unreduced_kind if unreduced else None\n    return mat.update(varying=varying, unreduced=unreduced, reduced=reduced,\n                      unreduced_kind=u_kind)\n  return mat\n\n\ndef get_memory_space(memory_space):\n  assert memory_space is not None\n  return memory_space\n\n\ndef _check_mat(varying, unreduced, reduced, unreduced_kind):\n  if varying & unreduced:\n    raise ValueError(\n        \"varying and unreduced cannot have common mesh axes. Got\"\n        f\" varying={varying} and unreduced={unreduced}\")\n  if varying & reduced:\n    raise ValueError(\n        \"varying and reduced cannot have common mesh axes. Got\"\n        f\" varying={varying} and reduced={reduced}\")\n  assert not (varying & unreduced & reduced)\n\n  if unreduced_kind is not None and not isinstance(unreduced_kind, UnreducedKind):\n    raise TypeError(\n        \"Expected unreduced_kind to be of type `jax.sharding.UnreducedKind`\"\n        f\" but got {type(unreduced_kind)}\")\n  if not unreduced and unreduced_kind is not None:\n    raise ValueError(\n        \"`unreduced_kind` should be `None` when `unreduced` is an empty set.\"\n        f\" Got {unreduced_kind=} and {unreduced=}\")\n\ndef _canonicalize_mat(name, val):\n  if not isinstance(val, frozenset):\n    if not isinstance(val, set):\n      raise TypeError(\n          f\"{name} argument of ManualAxisType should \"","sourceCodeStart":2331,"sourceCodeEnd":2367,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L2331-L2367","documentation":"_check_mat rejects a ManualAxisType where the same mesh axis is both 'varying' and 'reduced'. An axis is either replicated-varying or already reduced (result of a collective); it cannot be declared both.","triggerScenarios":"ManualAxisType(varying={'x'}, reduced={'x'}); typically from internal code or user-built avals after collectives like psum where the axis was already marked reduced.","commonSituations":"Composing collectives (psum then operations that assume varying) with manual mat bookkeeping; custom spmd primitives copying mat sets incorrectly.","solutions":["Remove the axis from 'reduced' if subsequent ops treat it as varying, or from 'varying' if the collective output is replicated","Derive mats automatically (JAX computes them for standard collectives) rather than setting them by hand"],"exampleFix":"// before\nmat = ManualAxisType(varying={'x'}, reduced={'x'})\n\n// after\nmat = ManualAxisType(varying={'x'})","handlingStrategy":"validation","validationCode":"if varying & reduced:\n    raise ValueError('axis in both varying and reduced')","typeGuard":"def disjoint_mat(varying, unreduced, reduced): return not (varying & unreduced) and not (varying & reduced)","tryCatchPattern":null,"preventionTips":["Compute reduced as full - varying - unreduced to guarantee disjointness"],"tags":["jax","sharding","manual-axis-type","validation"],"backgroundTag":"invalid-sharding-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}