jax-ml/jax · error · ValueError

reduced_vary_cast is a Reduced->Varying collective. This mea

Error message

reduced_vary_cast is a Reduced->Varying collective. This means that the axis names mentioned in `axes` passed to `reduced_vary_cast` must be present in `jax.typeof(x).mat.reduced`. Got axes={axes} and jax.typeof(x).mat.reduced={aval.mat.reduced}

What it means

For `reduced_vary_cast(x, axes)` every named axis in `axes` must be present in the input's `mat.reduced` set. The abstract eval raises when `aval.mat.reduced & set(axes)` is empty and reports both sets.

Source

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

batching.primitive_batchers[vary_unreduced_cast_p] = _vary_unreduced_cast_batcher

####################### reduced_vary_cast #############################

# Reduced -> Varying no-op cast
# Traceable defined in core.py to avoid circular imports
core.reduced_vary_cast_p.def_impl(
    partial(_raise_valueerror, 'reduced_vary_cast'))
mlir.register_lowering(core.reduced_vary_cast_p, lambda ctx, x, *, axes: [x])

def _reduced_vary_cast_abstract_eval(aval, *, axes):
  assert isinstance(axes, tuple)
  _check_axis_names(axes, 'reduced_vary_cast')
  if not aval.mat.reduced:
    raise ValueError('reduced_vary_cast only accepts inputs that are'
                     f' reduced. Got {aval.str_short(True)}')
  # If the intersection between aval.mat.reduced and axes is empty, error
  if not (aval.mat.reduced & set(axes)):
    raise ValueError(
        "reduced_vary_cast is a Reduced->Varying collective. This"
        " means that the axis names mentioned in `axes` passed to"
        " `reduced_vary_cast` must be present in"
        f" `jax.typeof(x).mat.reduced`. Got axes={axes} and"
        f" jax.typeof(x).mat.reduced={aval.mat.reduced}")
  if aval.mat.varying & set(axes):
    raise ValueError(
        "reduced_vary_cast input cannot be varying across the axis_name"
        f" provided. Got x={aval.str_short(True)} and axis_name={axes}")
  if aval.mat.unreduced:
    check_unreduced_kind('reduced_vary_cast', aval.mat, UnreducedKind.sum)

  new_reduced = frozenset(i for i in aval.mat.reduced if i not in axes)
  out_vma = aval.mat.varying | frozenset(axes)
  return aval.update(manual_axis_type=aval.mat.update(
    varying=out_vma, reduced=new_reduced))
core.reduced_vary_cast_p.def_abstract_eval(_reduced_vary_cast_abstract_eval)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check `jax.typeof(x).mat.reduced` and pass one of those axis names
2. Ensure the preceding psum/reduce collective used the same axis_name
3. Use `pcast` with the correct axis and let it dispatch

Example fix

// before
reduced_vary_cast(x, 'data')  # reduced over 'model'
// after
reduced_vary_cast(x, 'model')

t = jax.typeof(x)
assert set(axes) & t.mat.reduced, f'axes {axes} not reduced; reduced={t.mat.reduced}'

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

When it happens

Trigger: Casting to varying over axis 'i' when x is reduced over 'j'; using an axis name that was never reduced over (e.g. the value was reduced under a different mesh axis name).

Common situations: Renaming mesh axes or partition specs without updating cast sites; mixing collectives applied under different axis names in one pipeline.

Related errors

• vary_unreduced_cast is a Varying->Unreduced collective. This
• unbound axis name: {axis_name}
• unbound axis name: {axis_name}
• all_gather_reduced is a Varying -> Reduced collective. This
• unreduced_psum_scatter is a Unreduced -> Varying collective.

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