{"record":{"id":"6d4b54f6474ab31c","repo":"jax-ml/jax","slug":"pmap-in-axes-must-be-an-int-none-or-nested-con","errorCode":null,"errorMessage":"pmap in_axes must be an int, None, or (nested) container with those types as leaves, but got {in_axes}.","messagePattern":"pmap in_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":356,"sourceCode":"    )\n\n  wrapped.lower = lower  # pyrefly: ignore[missing-attribute]\n  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.","sourceCodeStart":338,"sourceCodeEnd":374,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pmap.py#L338-L374","documentation":"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.","triggerScenarios":"`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.","commonSituations":"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.","solutions":["Replace leaves with Python `int(...)` or None","Flatten/normalize numpy ints: `in_axes = tree_map(lambda a: None if a is None else int(a), in_axes)`","Verify each leaf is `type(l) is int` before calling pmap"],"exampleFix":"# before\nf = jax.pmap(fn, in_axes=(np.int64(0), None))\n# after\nf = jax.pmap(fn, in_axes=(int(np.int64(0)), None))","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(in_axes)), 'bad in_axes leaves'","typeGuard":"from jax.tree_util import tree_map\ndef normalize_axes(ax):\n    return tree_map(lambda a: None if a is None else int(a), ax)","tryCatchPattern":null,"preventionTips":["Coerce numpy ints to Python int","Keep axis specs as plain literals in config"],"tags":["jax","pmap","in-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"}