{"record":{"id":"13ad8a610ec3b3de","repo":"jax-ml/jax","slug":"preduced-input-cannot-be-reduced-across-the-axis-n","errorCode":null,"errorMessage":"preduced input cannot be reduced across the axis_name provided. Got x={aval.str_short(True)} and axis_name={axes}","messagePattern":"preduced input cannot be reduced across the axis_name provided\\. Got x=(.+?) and axis_name=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2886,"sourceCode":"  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)\n\ndef _preduced_batcher(vals_in, dims_in, *, axes):\n  raise NotImplementedError\nbatching.primitive_batchers[preduced_p] = _preduced_batcher\n\n######################## vary_unreduced_cast #######################","sourceCodeStart":2868,"sourceCodeEnd":2904,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2868-L2904","documentation":"`preduced` marks axes as reduced but refuses inputs that are already annotated as reduced across those axes (a double-reduction would be meaningless/incorrect). This ValueError fires when `aval.mat.reduced` intersects the requested axes.","triggerScenarios":"Calling `jax.lax.preduced(x, axis_name)` where `axis_name` is already in `jax.typeof(x).mat.reduced` — e.g. applying preduced twice for the same axis, or applying it to the output of a completed reduction.","commonSituations":"Autodiff transpose code applying preduced to values already marked reduced; accidentally chaining preduced with unreduced_psum outputs on the same axis; redundant manual annotations in custom collectives pipelines.","solutions":["Skip the preduced call for axes already reduced (guard with a set check on `mat.reduced`)","Restructure the pipeline so each axis is marked reduced exactly once","Inspect `jax.typeof(x).mat.reduced` before calling and pass only non-reduced axes"],"exampleFix":"# before\ny = jax.lax.preduced(jax.lax.preduced(x, 'data'), 'data')\n# after\ny = jax.lax.preduced(x, 'data')  # apply once per axis","handlingStrategy":"validation","validationCode":"import jax\nassert not (set(jax.typeof(x).mat.reduced) & set(axes)), 'axis already reduced'","typeGuard":"def axes_not_reduced(x, axes) -> bool:\n    return not (set(jax.typeof(x).mat.reduced) & set(axes))","tryCatchPattern":null,"preventionTips":["Apply preduced at most once per axis","Track reduced axes in your pipeline metadata","Guard transpose rules against double application"],"tags":["jax","sharding","collectives","double-reduction","mesh"],"backgroundTag":"sharding-annotation-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}