jax-ml/jax · error · ValueError

Got unexpected `to` value. Allowed `to` values are: {_allowe

Error message

Got unexpected `to` value. Allowed `to` values are: {_allowed_pcast_to}

What it means

`jax.lax.pcast`'s `to` keyword must be one of 'unreduced', 'reduced', or 'varying' (the set `_allowed_pcast_to`). Any other string raises this ValueError listing the allowed values.

Source

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

_pcast_funcs = {
    ('invarying', 'varying'): core.pvary,
    ('invarying', 'reduced'): preduced,
    ('varying', 'unreduced'): vary_unreduced_cast,
    ('reduced', 'varying'): core.reduced_vary_cast,
}

_allowed_pcast_to = {'unreduced', 'reduced', 'varying'}

def pcast(x, axis_name, *, to: str):
  if isinstance(axis_name, (set, frozenset)):
    raise TypeError(f"{axis_name=} must be a tuple or a str. Got {axis_name}")
  axes = (axis_name,) if not isinstance(axis_name, tuple) else axis_name
  if not axis_name:
    return x

  if to not in _allowed_pcast_to:
    raise ValueError(
        "Got unexpected `to` value. Allowed `to` values are:"
        f" {_allowed_pcast_to}")

  def bind(leaf):
    from_ = _get_from(core.typeof(leaf), axes, 'jax.lax.pcast')
    func = _pcast_funcs.get((from_, to), None)
    if func is None:
      raise ValueError(f"Unsupported pcast from={from_}, {to=}")
    return func(leaf, axes)
  return tree_util.tree_map(bind, x)

def _emit_async_start_custom_call(
    target_name, ctx, x, cfg, called_computations=None
):
  out_aval, = ctx.avals_out
  future_type = mlir.aval_to_ir_type(ctx.module_context, out_aval.inner_aval)

  cfg = dict(cfg)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use exactly one of 'unreduced', 'reduced', 'varying'
2. If targeting invarying, restructure: invarying is not a cast target — reshape/gather instead
3. Validate dynamic `to` values against the allowed set before calling

Example fix

# before
pcast(x, 'dev', to='invarying')
# after
# invarying is not a valid target; pick one of:
pcast(x, 'dev', to='varying')

ALLOWED = {'unreduced', 'reduced', 'varying'}
assert to in ALLOWED, f'to must be one of {ALLOWED}'

def is_valid_to(to: str) -> bool:
    return to in {'unreduced', 'reduced', 'varying'}

When it happens

Trigger: Calling `pcast(x, 'dev', to='invarying')`, `to='varying '` (typo/whitespace), or a placeholder like `to=None`/`to='auto'`.

Common situations: Assuming pcast can target every axis state including 'invarying'; typos or dynamically built `to` strings from config; older snippets using a different API vocabulary.

Related errors

• Unknown resize method "{s}"
• Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.
• 'trans' value must be 0, 1, or 2, got {trans}
• {name} can only accept axis_name which corresponds to one of
• {axis_name=} must be a tuple or a str. Got {axis_name}

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