jax-ml/jax · error · ValueError

vary_unreduced_cast is a Varying->Unreduced collective. This

Error message

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

What it means

For `vary_unreduced_cast(x, axes)` the axes you name must actually appear in the input's varying set. The abstract eval intersects `aval.mat.varying` with `set(axes)` and errors when the intersection is empty, telling you exactly which state the input had.

Source

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

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

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

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

  new_unreduced = aval.mat.unreduced | frozenset(axes)
  out_vma = frozenset(i for i in aval.mat.varying if i not in axes)
  return aval.update(manual_axis_type=aval.mat.update(
    varying=out_vma, unreduced=new_unreduced))
vary_unreduced_cast_p.def_abstract_eval(_vary_unreduced_cast_abstract_eval)

def _vary_unreduced_cast_transpose_rule(cts, x, *, axes):
  assert ad.is_undefined_primal(x)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Print `jax.typeof(x).mat.varying` and pass one of those exact axis names
2. Align the axis names with your `jax.sharding.Mesh` / NamedSharding axis names
3. Replace with `pcast(x, correct_axis, to='unreduced')`

Example fix

// before
jax.lax.vary_unreduced_cast(x, 'replica')  # x varies over 'data'
// after
jax.lax.vary_unreduced_cast(x, 'data')

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

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

When it happens

Trigger: Calling vary_unreduced_cast with axis_name 'i' when x is varying only over 'j' (or over a differently-named axis); typos or renamed axis names after refactoring.

Common situations: Renaming mesh/axis names in pjit sharding constraints or Mesh definitions without updating cast call sites; copying cast code between models with different axis naming conventions.

Related errors

• reduced_vary_cast is a Reduced->Varying collective. This mea
• 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/2c35e1d46f6873c4. Report an issue: GitHub.