{"record":{"id":"0da31f68c7e6e220","repo":"jax-ml/jax","slug":"axis-name-must-be-a-tuple-or-a-str-got-axis-n","errorCode":null,"errorMessage":"{axis_name=} must be a tuple or a str. Got {axis_name}","messagePattern":"(.+?) must be a tuple or a str\\. Got (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":3037,"sourceCode":"        f\"{name} can only accept axis_name which corresponds to one of\"\n        \" varying, unreduced, reduced or invarying state of the input. Got\"\n        f\" input type: {aval}, axes: {axes} and input state: {out}\")\n  o, = out\n  return o\n\n\n_pcast_funcs = {\n    ('invarying', 'varying'): core.pvary,\n    ('invarying', 'reduced'): preduced,\n    ('varying', 'unreduced'): vary_unreduced_cast,\n    ('reduced', 'varying'): core.reduced_vary_cast,\n}\n\n_allowed_pcast_to = {'unreduced', 'reduced', 'varying'}\n\ndef pcast(x, axis_name, *, to: str):\n  if isinstance(axis_name, (set, frozenset)):\n    raise TypeError(f\"{axis_name=} must be a tuple or a str. Got {axis_name}\")\n  axes = (axis_name,) if not isinstance(axis_name, tuple) else axis_name\n  if not axis_name:\n    return x\n\n  if to not in _allowed_pcast_to:\n    raise ValueError(\n        \"Got unexpected `to` value. Allowed `to` values are:\"\n        f\" {_allowed_pcast_to}\")\n\n  def bind(leaf):\n    from_ = _get_from(core.typeof(leaf), axes, 'jax.lax.pcast')\n    func = _pcast_funcs.get((from_, to), None)\n    if func is None:\n      raise ValueError(f\"Unsupported pcast from={from_}, {to=}\")\n    return func(leaf, axes)\n  return tree_util.tree_map(bind, x)\n\ndef _emit_async_start_custom_call(","sourceCodeStart":3019,"sourceCodeEnd":3055,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L3019-L3055","documentation":"`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.","triggerScenarios":"Passing `axis_name=frozenset(mesh.axis_names)` or a set comprehension of axis names to jax.lax.pcast.","commonSituations":"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.","solutions":["Convert to a sorted tuple: `tuple(sorted(names))`","Pass a single str when casting one axis","Keep axis names as tuples throughout your codebase, not sets"],"exampleFix":"# before\npcast(x, frozenset({'data','model'}), to='unreduced')\n# after\npcast(x, ('data','model'), to='unreduced')","handlingStrategy":"type-guard","validationCode":"if isinstance(axis_name, (set, frozenset)):\n    axis_name = tuple(sorted(axis_name))","typeGuard":"def is_str_or_tuple(v) -> bool:\n    return isinstance(v, str) or (isinstance(v, tuple) and all(isinstance(a, str) for a in v))","tryCatchPattern":null,"preventionTips":["Keep axis names as tuples or single strs throughout the codebase"],"tags":["jax","pcast","type-error","axis-name"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}