jax-ml/jax · error · ValueError

preduced is a Invariant->Reduced collective. This means that

Error message

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}

What it means

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

Source

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

  if not axes:
    return x
  cur_mesh = get_abstract_mesh()
  if not config._check_vma.value and all(a in cur_mesh.manual_axes for a in axes):
    return x
  new_axes = axes if cur_mesh.empty else core.order_wrt_mesh(cur_mesh, axes)
  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the axis from the input's varying set before calling `preduced` (e.g. via an appropriate reshard/annotation step)
2. Only pass axes along which the value is genuinely invariant
3. Check `jax.typeof(inp).mat.varying` first and reconcile the axes argument

Example fix

# before
jax.lax.preduced(x, 'data')  # x.varying contains 'data'
# after
jax.lax.preduced(x, 'data')  # after ensuring x is invariant along 'data'

import jax
assert not (set(jax.typeof(inp).mat.varying) & set(axes)), 'axis is varying; preduced needs invariant axes'

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

When it happens

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

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

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 input cannot be reduced across the axis_name provid

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