{"record":{"id":"4bef6a9a738fc7cf","repo":"jax-ml/jax","slug":"pmap-in-axes-value-not-a-valid-prefix-dynamic-me","errorCode":null,"errorMessage":"pmap in_axes: value not a valid prefix (dynamic message including 'The full pytree here is the tuple of arguments passed positionally to the pmapped function...' and 'Check that the value of the `in_axes` argument to `pmap` is a tree prefix of the tuple of arguments passed positionally to the pmapped function.')","messagePattern":"pmap in_axes: value not a valid prefix \\(dynamic message including 'The full pytree here is the tuple of arguments passed positionally to the pmapped function\\.\\.\\.' and 'Check that the value of the `in_axes` argument to `pmap` is a tree prefix of the tuple of arguments passed positionally to the pmapped function\\.'\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pmap.py","lineNumber":645,"sourceCode":"          \"\\n\\nThe 'full pytree' here is the tuple of arguments passed \"\n          \"positionally to the pmapped function, and the value of `in_axes` \"\n          \"must be a tree prefix of that tuple. But it was not a prefix.\"\n      )\n      if kwargs:\n        msg += (\n            \"\\n\\nWhen some arguments are passed by keyword to the pmapped \"\n            \"function, they are not included in the comparison to `in_axes`. \"\n            \"Instead, each argument passed by keyword is mapped over its \"\n            \"leading axis. See the description of `in_axes` in the `pmap` \"\n            \"docstring: \"\n            \"https://docs.jax.dev/en/latest/_autosummary/jax.pmap.html#jax.pmap\"\n        )\n      msg += (\n          \"\\n\\nCheck that the value of the `in_axes` argument to `pmap` \"\n          \"is a tree prefix of the tuple of arguments passed positionally to \"\n          \"the pmapped function.\"\n      )\n      raise ValueError(msg) from None\n\n  return in_axes_flat\n\n\ndef _get_donated_invars(donate_tuple, in_tree, num_flat_args):\n  \"\"\"Compute donation vector for arguments.\n\n  Args:\n    donate_tuple: Tuple of donated argument indices.\n    in_tree: PyTreeDef of input structure.\n    num_flat_args: Number of flat arguments.\n\n  Returns:\n    Tuple of bools indicating which flat args are donated.\n  \"\"\"\n\n  if donate_tuple and not config.debug_nans.value:\n    return donation_vector(donate_tuple, (), in_tree)","sourceCodeStart":627,"sourceCodeEnd":663,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pmap.py#L627-L663","documentation":"pmap flattens `in_axes` against the flattened tree of positional arguments; if in_axes is not a valid pytree prefix of the args tuple (wrong arity, wrong nesting, or extra levels), flattening fails and pmap raises a detailed prefix-mismatch ValueError.","triggerScenarios":"Calling a pmapped function with a different number/order of arguments than the in_axes tree describes, e.g. `in_axes=(0,0)` but calling with 3 args, or nested pytrees in args not matched by in_axes structure.","commonSituations":"Optional args passed sometimes and not others; refactored function signatures; passing kwargs (which are not part of the positional tuple) while in_axes assumed them positional.","solutions":["Make in_axes exactly mirror the structure of the positional arguments tuple (or use a scalar like 0)","Pass all mapped arguments positionally","Use `jax.tree_util.tree_structure(args)` to debug and match structures"],"exampleFix":"# before\nf = jax.pmap(fn, in_axes=(0, 0))\nf(x, y, z)\n# after\nf = jax.pmap(fn, in_axes=0)\nf(x, y, z)","handlingStrategy":"validation","validationCode":"from jax.tree_util import tree_flatten\nax_leaves = tree_flatten(in_axes)[0]\narg_leaves = tree_flatten(args)[0]\n# in_axes must be a prefix: flatten_axes would fail; quick arity sanity check\nassert len(ax_leaves) <= len(arg_leaves), 'in_axes arity exceeds args'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use scalar in_axes=0 unless per-arg mapping is needed","Test pmapped functions with the exact arg tuples used in production"],"tags":["jax","pmap","in-axes","pytree-prefix"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}