jax-ml/jax · error · TypeError

must be a tuple or a str. Got

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.

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: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• `addupdate` must be called on `Ref` types
• all_gather_reduced is a Varying -> Reduced collective. This…
• args argument to jax.scipy.optimize.minimize must be a…
• Argument ' ' of type ' ' is not a valid JAX type
• : arrays must have same number of dimensions, got .

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

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