{"record":{"id":"d5e9f60d7e3f35c8","repo":"jax-ml/jax","slug":"the-static-argnums-argument-to-jax-checkpoint","errorCode":null,"errorMessage":"the `static_argnums` argument to `jax.checkpoint` / `jax.remat` must be an int, tuple of ints or, bool, but got value {static_argnums}","messagePattern":"the `static_argnums` argument to `jax\\.checkpoint` / `jax\\.remat` must be an int, tuple of ints or, bool, but got value (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/ad_checkpoint.py","lineNumber":427,"sourceCode":"\n\ndef remat(fun: Callable, *, prevent_cse: bool = True,\n          policy: Callable[..., bool] | None = None,\n          static_argnums: int | tuple[int, ...] = ()) -> Callable:\n  \"\"\"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)","sourceCodeStart":409,"sourceCodeEnd":445,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/ad_checkpoint.py#L409-L445","documentation":"After rebuilding a structure, UnflattenImpl verifies the supplied leaf iterator was fully consumed; leftover leaves mean more leaves were passed than the treedef contains, raising 'Too many leaves for PyTreeDef; expected %d.'","triggerScenarios":"treedef.unflatten(leaves) / jax.tree.unflatten with len(leaves) > treedef.num_leaves: duplicating leaves (e.g. vstack of leaves), concatenating leaf lists from multiple structures, or unflattening a longer flat array without slicing to num_leaves.","commonSituations":"Broadcasting one pytree's leaves into another's structure; stacking params from several devices and unflattening the concatenated list; using leaves from a bigger tree with a smaller treedef after filtering the treedef instead of the leaves.","solutions":["Slice leaves to the expected count: treedef.unflatten(leaves[:treedef.num_leaves]) only after verifying that is semantically correct","Prefer jax.tree.unflatten over manual leaf list surgery; derive leaves with the same treedef","Validate len(leaves) == treedef.num_leaves before the call","Use jax.tree.map to redistribute values instead of concatenating leaf lists"],"exampleFix":"# before\nall_leaves = leaves_a + leaves_b\nout = treedef_a.unflatten(all_leaves)  # too many leaves\n\n# after\nout = treedef_a.unflatten(leaves_a)\nout_b = treedef_b.unflatten(leaves_b)","handlingStrategy":"validation","validationCode":"n = treedef.num_leaves\nassert len(leaves) == n, f'{len(leaves)} leaves supplied, treedef holds {n}'\nout = treedef.unflatten(leaves)","typeGuard":null,"tryCatchPattern":"try:\n    out = treedef.unflatten(leaves)\nexcept ValueError as e:\n    if 'Too many leaves' in str(e):\n        leaves = leaves[:treedef.num_leaves]  # only if truncation is intended\n        out = treedef.unflatten(leaves)\n    else:\n        raise","preventionTips":["Never concatenate leaf lists from different treedefs","Slice flat arrays to num_leaves (or use index_util) before unflatten","Prefer jax.tree.map over manual leaf-list plumbing"],"tags":["pytree","unflatten","jax","leaf-count"],"backgroundTag":"pytree-leaf-count-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}