jax-ml/jax · error · ValueError

vary_unreduced_cast only accepts inputs that are varying. Go

Error message

vary_unreduced_cast only accepts inputs that are varying. Got {aval.str_short(True)}

What it means

`vary_unreduced_cast` is a Varying->Unreduced no-op cast used in JAX's manual-parallelism (named axis) machinery. Its abstract eval requires the input aval to be marked as varying (`aval.mat.varying` non-empty). Passing a value that is not varying raises this ValueError.

Source

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

  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 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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check `jax.typeof(x).mat` before casting and only call vary_unreduced_cast on values whose `varying` set is non-empty
2. Use `jax.lax.pcast(x, axis_name, to='unreduced')`, which dispatches to the correct cast for the actual input state
3. Restructure the collective sequence so the operand is genuinely varying (e.g. re-materialize variation with a varying cast/split before)

Example fix

// before
y = jax.lax.vary_unreduced_cast(x, 'dev')  # x not varying
// after
y = jax.lax.pcast(x, 'dev', to='unreduced')

t = jax.typeof(x)
if not t.mat.varying:
    x = to_varying_state(x, axes)  # your normalization helper

def is_varying(x, ) -> bool:
    return bool(jax.typeof(x).mat.varying)

catch ValueError and re-dispatch via pcast(x, axes, to='unreduced')

When it happens

Trigger: Calling `jax.lax.vary_unreduced_cast(x, axis_name)` where `jax.typeof(x).mat.varying` is empty — e.g. the input was already fully reduced/unreduced/invarying across all named axes.

Common situations: Hand-writing cast sequences between manual axis states (varying/unreduced/reduced) and losing track of the current state; mixing pmap-style code with collectives that implicitly reduce, then casting the result.

Related errors

• reduced_vary_cast only accepts inputs that are reduced. Got
• reduced_vary_cast input cannot be varying across the axis_na
• vary_unreduced_cast is a Varying->Unreduced collective. This
• vary_unreduced_cast input cannot be unreduced across the axi
• reduced_vary_cast is a Reduced->Varying collective. This mea

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