{"record":{"id":"11c61873c100eba3","repo":"jax-ml/jax","slug":"no-batching-rule-defined-for-custom-vmap-function","errorCode":null,"errorMessage":"No batching rule defined for custom_vmap function {debug_fun.func_name} using def_vmap.","messagePattern":"No batching rule defined for custom_vmap function (.+?) using def_vmap\\.","errorType":"exception","errorClass":"AttributeError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_batching.py","lineNumber":156,"sourceCode":"    \"\"\"\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))\n    assert self.vmap_rule is not None\n    debug_rule = api_util.debug_info(\"custom_vmap rule\", self.vmap_rule,\n                                     (0, args, args), {})\n    out_flat = custom_vmap_p.bind(*consts, *args_flat,\n                                  call=closed_call,\n                                  rule=ClosedRule(self.vmap_rule,\n                                                  debug_rule),","sourceCodeStart":138,"sourceCodeEnd":174,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_batching.py#L138-L174","documentation":"A function decorated with @jax.custom_vmap must register a batching rule via <fun>.def_vmap(...) before it can be vmapped. Calling the wrapper without a registered vmap_rule raises this AttributeError.","triggerScenarios":"Decorating a function with @jax.custom_vmap but never calling fun.def_vmap(rule) and then invoking the function (directly or under jax.vmap).","commonSituations":"Forgetting the def_vmap step after adding the decorator; reordering code so def_vmap runs lazily/never (e.g. inside an uncalled branch); tutorial code that only shows the decorator half.","solutions":["Register a batching rule: call f.def_vmap with a rule mapping (axis_size, in_batched, *args) to (outs, out_batched)","Verify def_vmap executes at import time, not inside a conditional","Check the custom_vmap docstring example for the exact rule signature"],"exampleFix":"# before\n@jax.custom_vmap\ndef f(x): return x * 2\n# after\n@jax.custom_vmap\ndef f(x): return x * 2\n@f.def_vmap\ndef rule(axis_size, in_batched, x):\n    return x * 2, in_batched","handlingStrategy":"validation","validationCode":"assert getattr(f, 'vmap_rule', None) is not None, 'register f.def_vmap before vmapping'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Register def_vmap immediately below the decorated function","Add a jax.vmap smoke test at import time"],"tags":["jax","custom-vmap","vmap","missing-rule"],"backgroundTag":"missing-callback-registration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}