{"record":{"id":"d6f63dec8f5ca712","repo":"jax-ml/jax","slug":"varying-and-unreduced-cannot-have-common-mesh-axes","errorCode":null,"errorMessage":"varying and unreduced cannot have common mesh axes. Got varying={varying} and unreduced={unreduced}","messagePattern":"varying and unreduced cannot have common mesh axes\\. Got varying=(.+?) and unreduced=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":2345,"sourceCode":"  if config.remove_size_one_mesh_axis_from_type.value:\n    varying = frozenset(i for i in mat.varying\n                        if in_axis_env(i) or mesh.shape[i] != 1)\n    unreduced = frozenset(u for u in mat.unreduced if mesh.shape[u] != 1)\n    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):","sourceCodeStart":2327,"sourceCodeEnd":2363,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L2327-L2363","documentation":"_check_mat validates a ManualAxisType: a mesh axis cannot simultaneously be 'varying' (array differs across that axis) and 'unreduced' (kept un-reduced after a collective). The two states are mutually exclusive per axis.","triggerScenarios":"Building ManualAxisType(varying={'x'}, unreduced={'x'}) or updating an aval's mat so the same axis lands in both sets; usually via internal aval update_manual_axis_type calls or spmd primitives.","commonSituations":"Hand-constructed mats in custom collectives or shard_map plumbing; incorrect psum/ppermute wrappers tagging outputs unreduced on axes the input already varies on.","solutions":["Remove the axis from one of the two sets based on intended semantics","If the axis should vary, drop it from unreduced; if it must stay unreduced across a reduction, drop it from varying","Recompute mat from the primitive's output sharding instead of hand-writing it"],"exampleFix":"// before\nmat = ManualAxisType(varying={'x'}, unreduced={'x'})\n\n// after\nmat = ManualAxisType(varying={'x'})  # or unreduced={'x'} if semantically unreduced","handlingStrategy":"validation","validationCode":"if varying & unreduced:\n    raise ValueError('axis in both varying and unreduced')","typeGuard":"def disjoint_mat(varying, unreduced, reduced): return not (varying & unreduced) and not (varying & reduced)","tryCatchPattern":null,"preventionTips":["Always derive mat sets via set differences when composing","Unit-test mats for pairwise 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"}