jax-ml/jax · error · ValueError

preduced input cannot be reduced across the axis_name provid

Error message

preduced input cannot be reduced across the axis_name provided. Got x={aval.str_short(True)} and axis_name={axes}

What it means

`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.

Source

Thrown at jax/_src/lax/parallel.py:2886

  assert set(new_axes) == set(axes)
  del axes
  return tree_util.tree_map(lambda l: preduced_p.bind(l, axes=new_axes), x)

preduced_p = core.Primitive('preduced')
preduced_p.def_impl(partial(_raise_valueerror, 'preduced'))
mlir.register_lowering(preduced_p, lambda ctx, x, *, axes: [x])

def _preduced_abstract_eval(aval, *, axes):
  assert isinstance(axes, tuple)
  _check_axis_names(axes, 'preduced')
  if aval.mat.varying.intersection(set(axes)):
    raise ValueError(
        "preduced is a Invariant->Reduced collective. This means that the"
        " axis names mentioned in `axes` passed to `preduced` must not be"
        f" present in `jax.typeof(inp).mat.varying`. Got axes={axes} and"
        f" jax.typeof(inp).mat.varying={aval.mat.varying}")
  if aval.mat.reduced & set(axes):
    raise ValueError(
        "preduced input cannot be reduced across the axis_name"
        f" provided. Got x={aval.str_short(True)} and axis_name={axes}")
  if aval.mat.unreduced:
    check_unreduced_kind('preduced', aval.mat, UnreducedKind.sum)
  return aval.update(manual_axis_type=aval.mat.update(
      reduced=aval.mat.reduced | frozenset(axes)))
preduced_p.def_abstract_eval(_preduced_abstract_eval)

def _preduced_transpose_rule(cts, arg, *, axes):
  assert ad.is_undefined_primal(arg)
  return (unreduced_psum(cts, axis_name=axes),)
ad.deflinear2(preduced_p, _preduced_transpose_rule)

def _preduced_batcher(vals_in, dims_in, *, axes):
  raise NotImplementedError
batching.primitive_batchers[preduced_p] = _preduced_batcher

######################## vary_unreduced_cast #######################

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Skip the preduced call for axes already reduced (guard with a set check on `mat.reduced`)
2. Restructure the pipeline so each axis is marked reduced exactly once
3. Inspect `jax.typeof(x).mat.reduced` before calling and pass only non-reduced axes

Example fix

# before
y = jax.lax.preduced(jax.lax.preduced(x, 'data'), 'data')
# after
y = jax.lax.preduced(x, 'data')  # apply once per axis

import jax
assert not (set(jax.typeof(x).mat.reduced) & set(axes)), 'axis already reduced'

def axes_not_reduced(x, axes) -> bool:
    return not (set(jax.typeof(x).mat.reduced) & set(axes))

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• unbound axis name: {axis_name}
• {name} only accepts inputs that are unreduced. Got {aval.str
• {name} is a Unreduced -> Invariant collective. This means th
• {name}'s input cannot be varying across the axis_name provi
• preduced is a Invariant->Reduced collective. This means that

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/13ad8a610ec3b3de. Report an issue: GitHub.