{"record":{"id":"021d76e40196faab","repo":"jax-ml/jax","slug":"preduced-is-a-invariant-reduced-collective-this","errorCode":null,"errorMessage":"preduced is a Invariant->Reduced collective. This means that the axis names mentioned in `axes` passed to `preduced` must not be present in `jax.typeof(inp).mat.varying`. Got axes={axes} and jax.typeof(inp).mat.varying={aval.mat.varying}","messagePattern":"preduced is a Invariant->Reduced collective\\. This means that the axis names mentioned in `axes` passed to `preduced` must not be present in `jax\\.typeof\\(inp\\)\\.mat\\.varying`\\. Got axes=(.+?) and jax\\.typeof\\(inp\\)\\.mat\\.varying=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2880,"sourceCode":"  if not axes:\n    return x\n  cur_mesh = 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 core.order_wrt_mesh(cur_mesh, axes)\n  assert set(new_axes) == set(axes)\n  del axes\n  return tree_util.tree_map(lambda l: preduced_p.bind(l, axes=new_axes), x)\n\npreduced_p = core.Primitive('preduced')\npreduced_p.def_impl(partial(_raise_valueerror, 'preduced'))\nmlir.register_lowering(preduced_p, lambda ctx, x, *, axes: [x])\n\ndef _preduced_abstract_eval(aval, *, axes):\n  assert isinstance(axes, tuple)\n  _check_axis_names(axes, 'preduced')\n  if aval.mat.varying.intersection(set(axes)):\n    raise ValueError(\n        \"preduced is a Invariant->Reduced collective. This means that the\"\n        \" axis names mentioned in `axes` passed to `preduced` must not be\"\n        f\" present in `jax.typeof(inp).mat.varying`. Got axes={axes} and\"\n        f\" jax.typeof(inp).mat.varying={aval.mat.varying}\")\n  if aval.mat.reduced & set(axes):\n    raise ValueError(\n        \"preduced input cannot be reduced across the axis_name\"\n        f\" provided. Got x={aval.str_short(True)} and axis_name={axes}\")\n  if aval.mat.unreduced:\n    check_unreduced_kind('preduced', aval.mat, UnreducedKind.sum)\n  return aval.update(manual_axis_type=aval.mat.update(\n      reduced=aval.mat.reduced | frozenset(axes)))\npreduced_p.def_abstract_eval(_preduced_abstract_eval)\n\ndef _preduced_transpose_rule(cts, arg, *, axes):\n  assert ad.is_undefined_primal(arg)\n  return (unreduced_psum(cts, axis_name=axes),)\nad.deflinear2(preduced_p, _preduced_transpose_rule)","sourceCodeStart":2862,"sourceCodeEnd":2898,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2862-L2898","documentation":"`preduced` is an Invariant->Reduced no-op cast (the transpose of `unreduced_psum`). It requires that the axes you declare as reduced are not already marked as 'varying' in the input's materialization annotations. This ValueError fires when an axis in `axes` also appears in `jax.typeof(inp).mat.varying`.","triggerScenarios":"Calling `jax.lax.preduced(x, axis_name)` where the axis_name intersects `jax.typeof(x).mat.varying`, e.g. marking a value 'reduced' along an axis where it actually varies per device.","commonSituations":"Manual sharding-annotation bookkeeping in autodiff code (preduced appears in transposes of unreduced_psum); incorrect manual annotations after changing mesh axis semantics; mixing varying and reduced metadata on the same axis.","solutions":["Remove the axis from the input's varying set before calling `preduced` (e.g. via an appropriate reshard/annotation step)","Only pass axes along which the value is genuinely invariant","Check `jax.typeof(inp).mat.varying` first and reconcile the axes argument"],"exampleFix":"# before\njax.lax.preduced(x, 'data')  # x.varying contains 'data'\n# after\njax.lax.preduced(x, 'data')  # after ensuring x is invariant along 'data'","handlingStrategy":"validation","validationCode":"import jax\nassert not (set(jax.typeof(inp).mat.varying) & set(axes)), 'axis is varying; preduced needs invariant axes'","typeGuard":"def axes_invariant(x, axes) -> bool:\n    return not (set(jax.typeof(x).mat.varying) & set(axes))","tryCatchPattern":null,"preventionTips":["Check mat.varying before preduced in custom transpose rules","Maintain a single source of truth for axis annotations","Test custom collectives' forward+backward rules together"],"tags":["jax","sharding","collectives","transpose","mesh"],"backgroundTag":"sharding-annotation-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}