{"record":{"id":"8fa5e2dce53f8d8e","repo":"jax-ml/jax","slug":"pmap-requires-at-least-one-argument-with-a-mapped","errorCode":null,"errorMessage":"pmap requires at least one argument with a mapped axis.","messagePattern":"pmap requires at least one argument with a mapped axis\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pmap.py","lineNumber":498,"sourceCode":"    args: Flat list of arguments.\n    in_axes: Flat tuple of axis indices (int or None for each arg).\n\n  Returns:\n    The size of the mapped axis.\n\n  Raises:\n    ValueError: If no args have a mapped axis.\n  \"\"\"\n  if args and in_axes:\n    # Fast path: check first arg/axis (most common case).\n    if in_axes[0] is not None and hasattr(args[0], \"shape\"):\n      return int(args[0].shape[in_axes[0]])\n    # Slow path: scan for first mapped arg.\n    if isinstance(in_axes, tuple):\n      for arg, ax in zip(args, in_axes):\n        if ax is not None and hasattr(arg, \"shape\"):\n          return int(arg.shape[ax])\n  raise ValueError(\"pmap requires at least one argument with a mapped axis.\")\n\n\ndef _pmap_wrap_init(f, static_broadcasted_tuple):\n  \"\"\"Create a wrapped function with DebugInfo for pmap.\n\n  Args:\n    f: The function to wrap.\n    static_broadcasted_tuple: Tuple of static argument indices.\n\n  Returns:\n    A lu.WrappedFun ready for pmap.\n  \"\"\"\n  # Compute arg_names from signature, excluding static argnums\n  if (signature := fun_signature(f)) is not None:\n    static_set = frozenset(static_broadcasted_tuple)\n    arg_names = tuple(\n        name\n        for i, name in enumerate(signature.parameters.keys())","sourceCodeStart":480,"sourceCodeEnd":516,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pmap.py#L480-L516","documentation":"pmap infers the mapped axis size from arguments whose in_axes entry is not None. If every argument has in_axes=None (fully broadcast) and no explicit axis_size is available, there is no way to determine how many devices to map over.","triggerScenarios":"`jax.pmap(f, in_axes=None)(x)` where all leaves of in_axes are None, and no axis_size given; also when all args lack a `.shape` attribute (e.g. scalars/static objects).","commonSituations":"Broadcasting a constant computation over devices; passing static/metadata-only arguments; wrong in_axes tuple ordering so mapped axes align with non-array args.","solutions":["Set in_axes to an int (e.g. 0) for at least one array argument","Pass `axis_size=N` to pmap explicitly so the mapped size is known","Restructure to jit + sharding if no argument actually varies per device"],"exampleFix":"# before\nf = jax.pmap(fn, in_axes=None)\nf(x)\n# after\nf = jax.pmap(fn, in_axes=0)\nf(x)","handlingStrategy":"validation","validationCode":"from jax.tree_util import tree_leaves\nhas_mapped = any(a is not None for a in tree_leaves(in_axes))\nassert has_mapped or axis_size is not None, 'pmap needs a mapped axis or axis_size'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always give at least one arg in_axes=0","Pass axis_size explicitly when all args are broadcast"],"tags":["jax","pmap","in-axes","axis-size"],"backgroundTag":"unspecified-parallel-axis-size","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}