{"record":{"id":"f073f413fa88b105","repo":"jax-ml/jax","slug":"pmap-out-axes-must-be-an-int-none-or-nested-co","errorCode":null,"errorMessage":"pmap out_axes must be an int, None, or (nested) container with those types as leaves, but got {out_axes}.","messagePattern":"pmap out_axes must be an int, None, or \\(nested\\) container with those types as leaves, but got (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pmap.py","lineNumber":359,"sourceCode":"  return wrapped\n\n\ndef _prepare_pmap(fun, axis_name, static_broadcasted_argnums,\n                      donate_argnums, in_axes, out_axes):\n  # axis_size is an optional integer representing the global axis size.  The\n  # aggregate size (across all processes) size of the mapped axis must match the\n  # given value.\n  check_callable(fun)\n  axis_name = \"_internal_pmap_axis_name\" if axis_name is None else axis_name\n  static_broadcasted_tuple = _ensure_index_tuple(static_broadcasted_argnums)\n  donate_tuple = rebase_donate_argnums(\n      _ensure_index_tuple(donate_argnums), static_broadcasted_tuple)\n\n  if not all(type(l) is int for l in tree_leaves(in_axes)):\n    raise TypeError(\"pmap in_axes must be an int, None, or (nested) container \"\n                    f\"with those types as leaves, but got {in_axes}.\")\n  if not all(type(l) is int for l in tree_leaves(out_axes)):\n    raise TypeError(\"pmap out_axes must be an int, None, or (nested) container \"\n                    f\"with those types as leaves, but got {out_axes}.\")\n\n  return axis_name, static_broadcasted_tuple, donate_tuple\n\n\nclass CachedShardMap(NamedTuple):\n  \"\"\"Core cached pmap result.\n\n  Attributes:\n    pmapped: The shard_map-transformed function.\n    in_specs_flat: Flattened input PartitionSpecs for array conversion.\n    local_devices: List of devices in the local mesh.\n    in_local_shardings: NamedSharding for each input using local mesh.\n    in_global_shardings: NamedSharding for each input using global mesh.\n    mesh: The global Mesh for this pmap invocation.\n    out_specs: Output PartitionSpecs as a pytree prefix.\n    out_local_shardings_thunk: Cached thunk returning (local, global) sharding\n      pairs for output pspecs.","sourceCodeStart":341,"sourceCodeEnd":377,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pmap.py#L341-L377","documentation":"pmap's `out_axes` pytree leaves must be ints or None, mirroring in_axes for output mapping. Non-int leaves (floats, strings, numpy ints) fail the strict `type(l) is int` check in `_prepare_pmap`.","triggerScenarios":"`jax.pmap(f, out_axes=0.0)` or passing a nested out_axes spec containing non-int leaves, including numpy integer scalars.","commonSituations":"Reusing the same axes dict for in_axes/out_axes with mixed types; computed out_axes from config arrays; copy-paste from sharding specs.","solutions":["Coerce all leaves to Python int or None with tree_map","Check leaves with `type(a) is int` before calling pmap","Simplify: use scalar out_axes=0 if all outputs map on axis 0"],"exampleFix":"# before\nf = jax.pmap(fn, out_axes=(0, np.int32(0)))\n# after\nf = jax.pmap(fn, out_axes=0)","handlingStrategy":"type-guard","validationCode":"from jax.tree_util import tree_leaves\nassert all(type(a) is int or a is None for a in tree_leaves(out_axes)), 'bad out_axes leaves'","typeGuard":"def valid_axes_spec(ax) -> bool:\n    return ax is None or type(ax) is int or all(\n        type(a) is int or a is None for a in tree_leaves(ax))","tryCatchPattern":null,"preventionTips":["Prefer scalar out_axes=0","Run axis-spec lint in config loaders"],"tags":["jax","pmap","out-axes","type-validation"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}