jax-ml/jax · error · ValueError

pmapped function has static_broadcasted_argnums={static_broa

Error message

pmapped function has static_broadcasted_argnums={static_broadcasted_tuple} but was called with only {len(args)} positional argument{'s' if len(args) > 1 else ''}. All static broadcasted arguments must be passed positionally.

What it means

pmap's static_broadcasted_argnums index positional arguments. If the highest static index is >= the number of positional args actually passed at call time, the static indices point at nothing and pmap raises this ValueError.

Source

Thrown at jax/_src/pmap.py:544

  """Extract dynamic args and argnums after handling static args.

  Args:
    wrapped_f: The wrapped function.
    static_broadcasted_tuple: Tuple of static argument indices.
    args: Positional arguments.

  Returns:
    dyn_f: function with static args bound
    dyn_argnums: list of dynamic arg indices (or None if no static args)
    dyn_args: dynamic positional arguments (after static removed)

  Raises:
    ValueError: If static_broadcasted_argnums exceeds number of args.
  """

  if static_broadcasted_tuple:
    if max(static_broadcasted_tuple) >= len(args):
      raise ValueError(
          "pmapped function has"
          f" static_broadcasted_argnums={static_broadcasted_tuple} but was"
          f" called with only {len(args)} positional"
          f" argument{'s' if len(args) > 1 else ''}. All static broadcasted"
          " arguments must be passed positionally."
      )
    dyn_argnums = [
        i for i in range(len(args)) if i not in static_broadcasted_tuple
    ]
    wrapped_f, dyn_args = argnums_partial(wrapped_f, dyn_argnums, args)
  else:
    dyn_argnums = None
    dyn_args = args
  return wrapped_f, dyn_argnums, dyn_args


def _get_in_axes_flat(
    in_axes, dyn_argnums, dyn_args, kwargs, num_flat_args, in_tree

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass all static arguments positionally at every call site
2. Lower static_broadcasted_argnums indices to match the actual signature
3. Move the value into a closure or functools.partial instead of static argnums

Example fix

# before
f_pmapped = jax.pmap(f, static_broadcasted_argnums=1)
f_pmapped(x)  # missing second positional arg
# after
f_pmapped(x, n)  # pass static arg positionally

n_static = len(args)
assert all(i < n_static for i in static_broadcasted_tuple), 'static argnum out of range'

try:
    f_pmapped(*args)
except ValueError as e:
    if 'static_broadcasted_argnums' in str(e):
        raise  # fix call site: pass static args positionally
    raise

When it happens

Trigger: `jax.pmap(f, static_broadcasted_argnums=1)` but calling `f(x)` with one positional arg; or later adding static argnums without updating call sites; passing the static value by keyword instead of positionally.

Common situations: Refactoring functions to add parameters; passing hyperparameters as kwargs (which pmap cannot count); varying call signatures across callers.

Related errors

• {_vjp_too_many_args(name, len(extra_args) + 1)}
• unbound axis name: {axis_name}
• Cannot lower jaxpr with effects: {closed_jaxpr.effects}
• scan number of arguments doesn't match the number of jaxpr a
• Expected num_classes to match the size of axis {axis}, but {

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