jax-ml/jax · error · TypeError

pmap in_axes must be an int, None, or (nested) container wit

Error message

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

What it means

pmap's `in_axes` pytree leaves must be ints or None, indicating which axis of each argument is mapped. Any other leaf type (float, string, list, etc.) is rejected during pmap setup.

Source

Thrown at jax/_src/pmap.py:356

    )

  wrapped.lower = lower  # pyrefly: ignore[missing-attribute]
  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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace leaves with Python `int(...)` or None
2. Flatten/normalize numpy ints: `in_axes = tree_map(lambda a: None if a is None else int(a), in_axes)`
3. Verify each leaf is `type(l) is int` before calling pmap

Example fix

# before
f = jax.pmap(fn, in_axes=(np.int64(0), None))
# after
f = jax.pmap(fn, in_axes=(int(np.int64(0)), None))

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

from jax.tree_util import tree_map
def normalize_axes(ax):
    return tree_map(lambda a: None if a is None else int(a), ax)

When it happens

Trigger: `jax.pmap(f, in_axes=0.0)`, `in_axes=(0, 'batch')`, or a nested container containing a non-int leaf; also NumPy integer types fail the strict `type(l) is int` check.

Common situations: Using np.int64 values from computed configs; confusing pmap in_axes with vjp/vmap-style axis specs or with shard_map specs; typos like True instead of 0.

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 out_axes must be an int, None, or (nested) container wi
• pmap requires at least one argument with a mapped axis.
• pmap in_axes: value not a valid prefix (dynamic message incl
• {name} was requested to map a value of non-array type {core.
• {name} was requested to map its argument along axis {axis},

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