jax-ml/jax · error · TypeError

{axis_name=} must be a tuple or a str. Got {axis_name}

Error message

{axis_name=} must be a tuple or a str. Got {axis_name}

What it means

`jax.lax.pcast` explicitly rejects set/frozenset values for `axis_name`, requiring a str or a tuple of strs. The check exists because set ordering is nondeterministic, which would make cast semantics ambiguous.

Source

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

        f"{name} can only accept axis_name which corresponds to one of"
        " varying, unreduced, reduced or invarying state of the input. Got"
        f" input type: {aval}, axes: {axes} and input state: {out}")
  o, = out
  return o


_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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert to a sorted tuple: `tuple(sorted(names))`
2. Pass a single str when casting one axis
3. Keep axis names as tuples throughout your codebase, not sets

Example fix

# before
pcast(x, frozenset({'data','model'}), to='unreduced')
# after
pcast(x, ('data','model'), to='unreduced')

if isinstance(axis_name, (set, frozenset)):
    axis_name = tuple(sorted(axis_name))

def is_str_or_tuple(v) -> bool:
    return isinstance(v, str) or (isinstance(v, tuple) and all(isinstance(a, str) for a in v))

When it happens

Trigger: Passing `axis_name=frozenset(mesh.axis_names)` or a set comprehension of axis names to jax.lax.pcast.

Common situations: Programmatically deriving axis names from a Mesh's `axis_names` (which is a tuple but often converted to set) and passing them through; sharing helper code where other JAX APIs accepted sets.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li
• {prim_name} takes a scalar pred as argument, got {pred}
• `compute_on`'s compute_type argument must be a string.
• unbound axis name: {axis_name}

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