{"record":{"id":"fd19d30e63f56051","repo":"jax-ml/jax","slug":"the-static-argnums-argument-to-jax-checkpoint-fd19d3","errorCode":null,"errorMessage":"the `static_argnums` argument to `jax.checkpoint` / `jax.remat` can only take integer values greater than or equal to `-len(args)` and less than `len(args)`, but got {static_argnums}, while `len(args)` = {len(args)}","messagePattern":"the `static_argnums` argument to `jax\\.checkpoint` / `jax\\.remat` can only take integer values greater than or equal to `-len\\(args\\)` and less than `len\\(args\\)`, but got (.+?), while `len\\(args\\)` = (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/ad_checkpoint.py","lineNumber":432,"sourceCode":"  \"\"\"Alias of :func:`jax.checkpoint`.\"\"\"\n  return checkpoint(fun, prevent_cse=prevent_cse, policy=policy,\n                    static_argnums=static_argnums)\n\n# This function is similar to api_util.argnums_partial, except the error\n# messages are specific to jax.remat (and thus more actionable), the\n# hashing/caching behavior is slightly different, and this function accepts a\n# boolean for static_argnums. Perhaps the two could be de-duplicated.\ndef _remat_static_argnums(fun, static_argnums, args):\n  if type(static_argnums) is int:\n    static_argnums = (static_argnums,)\n  elif not (type(static_argnums) is tuple and\n            all(type(d) is int for d in static_argnums)):\n    raise TypeError(\"the `static_argnums` argument to `jax.checkpoint` / \"\n                    \"`jax.remat` must be an int, tuple of ints or, bool, but \"\n                    f\"got value {static_argnums}\")\n\n  if not all(-len(args) <= d < len(args) for d in static_argnums):\n    raise ValueError(\"the `static_argnums` argument to `jax.checkpoint` / \"\n                     \"`jax.remat` can only take integer values greater than or \"\n                     \"equal to `-len(args)` and less than `len(args)`, but got \"\n                     f\"{static_argnums}, while `len(args)` = {len(args)}\")\n\n  if not static_argnums:\n    return fun, args\n  nargs = len(args)\n  static_argnums_ = frozenset(d % len(args) for d in static_argnums)\n  dyn_args, static_args = [], []\n  for i, x in enumerate(args):\n    if i in static_argnums_: static_args.append(WrapHashably(x))\n    else: dyn_args.append(x)\n  new_fun = _dyn_args_fun(fun, static_argnums_, tuple(static_args), nargs)\n  return new_fun, dyn_args\n\nWrapHashably = api_util.WrapHashably\n_dyn_args_fun = api_util.dyn_args_fun\n","sourceCodeStart":414,"sourceCodeEnd":450,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/ad_checkpoint.py#L414-L450","documentation":"FlattenUpTo walks the receiver treedef's traversal over a user-supplied tree xs; if xs runs out of nodes while the treedef still expects structure (agenda empty but traversal unfinished), the prefixes don't match and jaxlib throws 'Tree structures did not match: %s vs %s' with the repr of xs and the treedef string.","triggerScenarios":"Calling treedef.flatten_up_to(xs) (used internally by jax.tree_util.tree_map when mapping one structure over another with a different prefix, e.g. tree_map(f, tree, smaller_tree) or tree_map(f, params, batched_shorthand)) where xs lacks a node the treedef requires — fewer dict keys, shorter tuple/list, or a leaf where a container was expected.","commonSituations":"tree_map over two dicts with mismatched keys; applying a single-element tree over a multi-element structure; default-device-array or config object passed where a pytree of arrays is expected; shape/batching bugs where the second argument lost a nesting level.","solutions":["Compare structures first: jax.tree.structure(a) == jax.tree.structure(b), or tree_map(print, ...) to inspect","Make the arguments' prefixes match (e.g. add missing dict keys, fix nesting depth)","If mapping a constant/tree over another, put the tree with the larger/equal structure first: tree_map(f, big, small)","Use is_leaf or jax.tree.map with None handling to explicitly define prefix semantics"],"exampleFix":"# before\nparams = {'w': ..., 'b': ...}\nstats = {'w': ...}\njax.tree.map(f, params, stats)  # structures did not match\n\n# after\nstats = {'w': ..., 'b': zeros_like(params['b'])}\njax.tree.map(f, params, stats)","handlingStrategy":"validation","validationCode":"def prefixes_match(treedef, xs) -> bool:\n    try:\n        treedef.flatten_up_to(xs)\n        return True\n    except ValueError:\n        return False\n\n# or:\n# jax.tree.structure(a) == jax.tree.structure(b)","typeGuard":"def same_structure(a, b) -> bool:\n    return jax.tree.structure(a) == jax.tree.structure(b)","tryCatchPattern":"try:\n    jax.tree.map(f, a, b)\nexcept ValueError as e:\n    if 'Tree structures did not match' in str(e):\n        raise ValueError(f'structure mismatch:\\n{jax.tree.structure(a)}\\nvs\\n{jax.tree.structure(b)}') from e\n    raise","preventionTips":["Check jax.tree.structure equality before tree_map over two trees","Put the structurally-larger tree first in multi-argument tree_map","Keep dict keys in sync (e.g. derive one tree from the other via tree_map)"],"tags":["pytree","tree-map","structure-mismatch","jax"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}