jax-ml/jax · error · ValueError

pvary is a invariant->variant collective. This means that th

Error message

pvary is a invariant->variant collective. This means that the axis names mentioned in `axes` passed to `pvary` must not be present in `jax.typeof(inp).mat.varying`. Got axes={axes} and jax.typeof(inp)={aval}

What it means

`pvary` marks an array as varying along named mesh axes; it is an invariant->variant collective, so JAX raises this error when any axis in `axes` is already present in the array's `mat.varying` set. Re-varying an already-varying axis is redundant and usually indicates a misunderstanding of the mesh annotation or a duplicated call. The abstract eval in `_pvary_abstract_eval` rejects it before any computation runs.

Source

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

def _psum_invariant_transpose_rule(cts, arg, *, axes):
  assert ad.is_undefined_primal(arg)
  return (core.pvary(cts, axis_name=axes),)
ad.deflinear2(psum_invariant_p, _psum_invariant_transpose_rule)

########################### pvary ##################################

def _raise_valueerror(name, arg, *, axes):
  raise ValueError(f'{name} should be called under jax.shard_map.')

core.pvary_p.def_impl(partial(_raise_valueerror, 'pvary'))
mlir.register_lowering(core.pvary_p, lambda ctx, x, *, axes: [x])

def _pvary_abstract_eval(aval, *, axes):
  _check_axis_names(axes, 'pvary')
  check_unreduced_args([aval], axes, 'pvary')
  assert isinstance(axes, tuple)
  if set(axes).intersection(aval.mat.varying):
    raise ValueError(
        "pvary is a invariant->variant collective. This means that the axis"
        " names mentioned in `axes` passed to `pvary` must not be present in"
        f" `jax.typeof(inp).mat.varying`. Got axes={axes} and"
        f" jax.typeof(inp)={aval}")
  out_vma = aval.mat.varying.union(frozenset(axes))
  return aval.update(sharding=aval.sharding.update(mesh=get_abstract_mesh()),
                     manual_axis_type=aval.mat.update(varying=out_vma))
core.pvary_p.def_abstract_eval(_pvary_abstract_eval)

def _pvary_transpose_rule(cts, arg, *, axes):
  assert ad.is_undefined_primal(arg)
  return (psum_invariant_p.bind(cts, axes=axes),)
ad.deflinear2(core.pvary_p, _pvary_transpose_rule)

def _pvary_batcher(vals_in, dims_in, *, axes):
  if any(type(axis) is int for axis in axes):
    raise NotImplementedError
  (x,), (d,) = vals_in, dims_in

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the redundant `pvary` call for axes already listed in `jax.typeof(inp).mat.varying`
2. Inspect `jax.typeof(x).mat.varying` before calling pvary and only pass axes not present
3. Check upstream collectives (e.g. a previous pvary or all_gather) that may already vary the axis
4. If you intended a different semantic, use the appropriate collective (e.g. pbroadcast for the inverse direction)

Example fix

// before
y = lax.pvary(lax.pvary(x, 'i'), 'i')
// after
y = lax.pvary(x, 'i')

varying = jax.typeof(x).mat.varying
new_axes = tuple(a for a in axes if a not in varying)
y = lax.pvary(x, new_axes) if new_axes else x

def can_pvary(x, axes: tuple[str, ...]) -> bool:
    return not (set(axes) & jax.typeof(x).mat.varying)

try:
    y = lax.pvary(x, axes)
except ValueError as e:
    if 'pvary is a invariant->variant' in str(e):
        y = x  # already varying on those axes
    else:
        raise

When it happens

Trigger: Calling `lax.pvary(x, 'i')` twice with the same axis name; or passing an axis name that is already in `jax.typeof(x).mat.varying` because a previous pvary/pvarying operation or a collective produced it.

Common situations: Refactoring older `pmap`/axis-index code to the new named-mesh `pvary` API and accidentally applying pvary after another primitive that already varied that axis; stacking collectives in a pipeline where the output is already varying.

Related errors

• all_gather_reduced only accepts inputs that are varying. Got
• unreduced_psum_scatter only accepts inputs that are unreduce
• callbacks are only supported in spmd computations when all m
• {name} cannot accept args which are unreduced. Got {a.str_sh
• {name} cannot accept args with unreduced_kind={a.mat.unreduc

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