jax-ml/jax · error · ValueError

the `static_argnums` argument to `jax.checkpoint` / `jax.rem

Error message

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)}

What it means

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.

Source

Thrown at jax/_src/ad_checkpoint.py:432

  """Alias of :func:`jax.checkpoint`."""
  return checkpoint(fun, prevent_cse=prevent_cse, policy=policy,
                    static_argnums=static_argnums)

# This function is similar to api_util.argnums_partial, except the error
# messages are specific to jax.remat (and thus more actionable), the
# hashing/caching behavior is slightly different, and this function accepts a
# boolean for static_argnums. Perhaps the two could be de-duplicated.
def _remat_static_argnums(fun, static_argnums, args):
  if type(static_argnums) is int:
    static_argnums = (static_argnums,)
  elif not (type(static_argnums) is tuple and
            all(type(d) is int for d in static_argnums)):
    raise TypeError("the `static_argnums` argument to `jax.checkpoint` / "
                    "`jax.remat` must be an int, tuple of ints or, bool, but "
                    f"got value {static_argnums}")

  if not all(-len(args) <= d < len(args) for d in static_argnums):
    raise ValueError("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 "
                     f"{static_argnums}, while `len(args)` = {len(args)}")

  if not static_argnums:
    return fun, args
  nargs = len(args)
  static_argnums_ = frozenset(d % len(args) for d in static_argnums)
  dyn_args, static_args = [], []
  for i, x in enumerate(args):
    if i in static_argnums_: static_args.append(WrapHashably(x))
    else: dyn_args.append(x)
  new_fun = _dyn_args_fun(fun, static_argnums_, tuple(static_args), nargs)
  return new_fun, dyn_args

WrapHashably = api_util.WrapHashably
_dyn_args_fun = api_util.dyn_args_fun

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Compare structures first: jax.tree.structure(a) == jax.tree.structure(b), or tree_map(print, ...) to inspect
2. Make the arguments' prefixes match (e.g. add missing dict keys, fix nesting depth)
3. If mapping a constant/tree over another, put the tree with the larger/equal structure first: tree_map(f, big, small)
4. Use is_leaf or jax.tree.map with None handling to explicitly define prefix semantics

Example fix

# before
params = {'w': ..., 'b': ...}
stats = {'w': ...}
jax.tree.map(f, params, stats)  # structures did not match

# after
stats = {'w': ..., 'b': zeros_like(params['b'])}
jax.tree.map(f, params, stats)

def prefixes_match(treedef, xs) -> bool:
    try:
        treedef.flatten_up_to(xs)
        return True
    except ValueError:
        return False

# or:
# jax.tree.structure(a) == jax.tree.structure(b)

def same_structure(a, b) -> bool:
    return jax.tree.structure(a) == jax.tree.structure(b)

try:
    jax.tree.map(f, a, b)
except ValueError as e:
    if 'Tree structures did not match' in str(e):
        raise ValueError(f'structure mismatch:\n{jax.tree.structure(a)}\nvs\n{jax.tree.structure(b)}') from e
    raise

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Effects not supported in AD of `checkpoint`/`remat`: {disall
• non-bool prevent_cse together with static_argnums/static_arg
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/fd19d30e63f56051. Report an issue: GitHub.