{"record":{"id":"3337eb7702cbce22","repo":"jax-ml/jax","slug":"the-input-arguments-to-the-custom-vmap-decorated-f","errorCode":null,"errorMessage":"The input arguments to the custom_vmap-decorated function {debug_fun.func_name} could not be resolved to positional-only arguments. Binding failed with the error:\\n{e}","messagePattern":"The input arguments to the custom_vmap-decorated function (.+?) could not be resolved to positional-only arguments\\. Binding failed with the error:\\\\n(.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_batching.py","lineNumber":149,"sourceCode":"        and (3) the batched arguments. It should return a tuple of the batched\n        output and a pytree of booleans with the same structure as the output,\n        specifying whether each output element is batched. See the documentation\n        for :py:func:`jax.custom_batching.custom_vmap` for some examples.\n\n    Returns:\n      This method passes the rule through, returning ``vmap_rule`` unchanged.\n    \"\"\"\n    self.vmap_rule = vmap_rule\n    return vmap_rule\n\n  @traceback_util.api_boundary\n  def __call__(self, *args, **kwargs):\n    debug_fun = api_util.debug_info(\"custom_vmap fun\", self.fun,\n                                    args, kwargs)\n    try:\n      args = api_util.resolve_kwargs(self.fun, args, kwargs)\n    except TypeError as e:\n      raise TypeError(\n          \"The input arguments to the custom_vmap-decorated function \"\n          f\"{debug_fun.func_name} could not be resolved to positional-only \"\n          f\"arguments. Binding failed with the error:\\n{e}\"\n      ) from e\n\n    if not self.vmap_rule:\n      raise AttributeError(\n          f\"No batching rule defined for custom_vmap function {debug_fun.func_name} \"\n          \"using def_vmap.\")\n    args_flat, in_tree = tracing_registry.flatten(args)\n    in_avals = [core.typeof(x) for x in args_flat]\n    jaxpr, out_avals = pe.trace_to_jaxpr(\n        self.fun, ft.pack((ft.treedef_args_to_ft(in_tree, in_avals), {})),\n        debug_fun)\n    closed_call, consts = pe.separate_consts(jaxpr)\n    out_tree = tree_structure(out_avals.unflatten())\n    in_tree = treedef_tuple_tracing_registry(\n        (tracing_registry.flatten(consts)[1], in_tree))","sourceCodeStart":131,"sourceCodeEnd":167,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_batching.py#L131-L167","documentation":"jax.custom_vmap-decorated functions must be callable with positional arguments only. This error wraps the TypeError raised when Python cannot bind the supplied keyword arguments to the function signature (unknown kwarg, or kwargs that the flattener cannot reorder).","triggerScenarios":"Calling a @jax.custom_vmap function with keyword arguments, e.g. f(x, scale=2.0) where scale is a regular parameter, or passing a kwarg name not present in the signature.","commonSituations":"Refactoring a function signature (renaming params) while callers still use old kwarg names; applying vmap/jit pipelines that forward **kwargs into custom_vmap functions.","solutions":["Call the function with all arguments positionally","If keyword usage is desired, wrap the function to bind kwargs first or accept **kwargs explicitly and delegate positionally"],"exampleFix":"# before\nout = f(x, y, scale=2.0)\n# after\nout = f(x, y, 2.0)","handlingStrategy":"validation","validationCode":"import inspect\ndef call_positional(fn, *args, **kwargs):\n    sig = inspect.signature(fn)\n    return fn(*sig.bind(None, *args, **kwargs).args) if kwargs else fn(*args)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use positional-only markers (/) in custom_vmap signatures to surface errors at def time","Lint for kwarg calls into custom_vmap functions"],"tags":["jax","custom-vmap","kwargs","argument-binding"],"backgroundTag":"keyword-argument-binding-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}