jax-ml/jax · error · TypeError

pmap out_axes must be an int, None, or (nested) container wi

Error message

pmap out_axes must be an int, None, or (nested) container with those types as leaves, but got {out_axes}.

What it means

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`.

Source

Thrown at jax/_src/pmap.py:359

  return wrapped


def _prepare_pmap(fun, axis_name, static_broadcasted_argnums,
                      donate_argnums, in_axes, out_axes):
  # axis_size is an optional integer representing the global axis size.  The
  # aggregate size (across all processes) size of the mapped axis must match the
  # given value.
  check_callable(fun)
  axis_name = "_internal_pmap_axis_name" if axis_name is None else axis_name
  static_broadcasted_tuple = _ensure_index_tuple(static_broadcasted_argnums)
  donate_tuple = rebase_donate_argnums(
      _ensure_index_tuple(donate_argnums), static_broadcasted_tuple)

  if not all(type(l) is int for l in tree_leaves(in_axes)):
    raise TypeError("pmap in_axes must be an int, None, or (nested) container "
                    f"with those types as leaves, but got {in_axes}.")
  if not all(type(l) is int for l in tree_leaves(out_axes)):
    raise TypeError("pmap out_axes must be an int, None, or (nested) container "
                    f"with those types as leaves, but got {out_axes}.")

  return axis_name, static_broadcasted_tuple, donate_tuple


class CachedShardMap(NamedTuple):
  """Core cached pmap result.

  Attributes:
    pmapped: The shard_map-transformed function.
    in_specs_flat: Flattened input PartitionSpecs for array conversion.
    local_devices: List of devices in the local mesh.
    in_local_shardings: NamedSharding for each input using local mesh.
    in_global_shardings: NamedSharding for each input using global mesh.
    mesh: The global Mesh for this pmap invocation.
    out_specs: Output PartitionSpecs as a pytree prefix.
    out_local_shardings_thunk: Cached thunk returning (local, global) sharding
      pairs for output pspecs.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Coerce all leaves to Python int or None with tree_map
2. Check leaves with `type(a) is int` before calling pmap
3. Simplify: use scalar out_axes=0 if all outputs map on axis 0

Example fix

# before
f = jax.pmap(fn, out_axes=(0, np.int32(0)))
# after
f = jax.pmap(fn, out_axes=0)

from jax.tree_util import tree_leaves
assert all(type(a) is int or a is None for a in tree_leaves(out_axes)), 'bad out_axes leaves'

def valid_axes_spec(ax) -> bool:
    return ax is None or type(ax) is int or all(
        type(a) is int or a is None for a in tree_leaves(ax))

When it happens

Trigger: `jax.pmap(f, out_axes=0.0)` or passing a nested out_axes spec containing non-int leaves, including numpy integer scalars.

Common situations: Reusing the same axes dict for in_axes/out_axes with mixed types; computed out_axes from config arrays; copy-paste from sharding specs.

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

• pmap in_axes must be an int, None, or (nested) container wit
• unbound axis name: {axis_name}
• Argument '{arg}' of type {type(arg)} is not a valid JAX type
• SymbolicScope constraint must be a string: got {repr(c_str)}
• Custom VJP bwd rule {self.bwd} was registered with defvjp_wi

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