{"record":{"id":"5588eeb4e5d98992","repo":"jax-ml/jax","slug":"pmapped-function-has-static-broadcasted-argnums-s","errorCode":null,"errorMessage":"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.","messagePattern":"pmapped function has static_broadcasted_argnums=(.+?) but was called with only (.+?) positional argument(.+?)\\. All static broadcasted arguments must be passed positionally\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pmap.py","lineNumber":544,"sourceCode":"  \"\"\"Extract dynamic args and argnums after handling static args.\n\n  Args:\n    wrapped_f: The wrapped function.\n    static_broadcasted_tuple: Tuple of static argument indices.\n    args: Positional arguments.\n\n  Returns:\n    dyn_f: function with static args bound\n    dyn_argnums: list of dynamic arg indices (or None if no static args)\n    dyn_args: dynamic positional arguments (after static removed)\n\n  Raises:\n    ValueError: If static_broadcasted_argnums exceeds number of args.\n  \"\"\"\n\n  if static_broadcasted_tuple:\n    if max(static_broadcasted_tuple) >= len(args):\n      raise ValueError(\n          \"pmapped function has\"\n          f\" static_broadcasted_argnums={static_broadcasted_tuple} but was\"\n          f\" called with only {len(args)} positional\"\n          f\" argument{'s' if len(args) > 1 else ''}. All static broadcasted\"\n          \" arguments must be passed positionally.\"\n      )\n    dyn_argnums = [\n        i for i in range(len(args)) if i not in static_broadcasted_tuple\n    ]\n    wrapped_f, dyn_args = argnums_partial(wrapped_f, dyn_argnums, args)\n  else:\n    dyn_argnums = None\n    dyn_args = args\n  return wrapped_f, dyn_argnums, dyn_args\n\n\ndef _get_in_axes_flat(\n    in_axes, dyn_argnums, dyn_args, kwargs, num_flat_args, in_tree","sourceCodeStart":526,"sourceCodeEnd":562,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pmap.py#L526-L562","documentation":"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.","triggerScenarios":"`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.","commonSituations":"Refactoring functions to add parameters; passing hyperparameters as kwargs (which pmap cannot count); varying call signatures across callers.","solutions":["Pass all static arguments positionally at every call site","Lower static_broadcasted_argnums indices to match the actual signature","Move the value into a closure or functools.partial instead of static argnums"],"exampleFix":"# before\nf_pmapped = jax.pmap(f, static_broadcasted_argnums=1)\nf_pmapped(x)  # missing second positional arg\n# after\nf_pmapped(x, n)  # pass static arg positionally","handlingStrategy":"validation","validationCode":"n_static = len(args)\nassert all(i < n_static for i in static_broadcasted_tuple), 'static argnum out of range'","typeGuard":null,"tryCatchPattern":"try:\n    f_pmapped(*args)\nexcept ValueError as e:\n    if 'static_broadcasted_argnums' in str(e):\n        raise  # fix call site: pass static args positionally\n    raise","preventionTips":["Pass static args positionally, never by keyword","Re-audit call sites after changing static_broadcasted_argnums"],"tags":["jax","pmap","static-argnums","argument-count"],"backgroundTag":"missing-positional-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}