code: def __call__(self, *args, **kwargs):
      assert not kwargs
      args_flat, in_tree = tree_flatten(args)
      flat_fun, out_tree = flatten_fun_nokwargs(
          lu.wrap_init(self.fun), in_tree
      )
      in_avals = [
          core.raise_to_shaped(core.get_aval(x)) for x in args_flat
      ]
      debug = pe.debug_info(
          self.fun, in_tree, out_tree, False, "custom_vmap"
      )
      jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(
          flat_fun, in_avals, debug
      )
      closed_call = core.ClosedJaxpr(
          pe.convert_constvars_jaxpr(jaxpr), ()
      )
      in_tree = treedef_tuple((tree_structure(consts), in_tree))
      assert self.vmap_rule is not None
      out_flat = custom_vmap_p.bind(
          *consts,
          *args_flat,
          call=closed_call,
          rule=ClosedRule(self.vmap_rule),
          in_tree=in_tree,
          out_tree=out_tree(),
      )
      return tree_unflatten(out_tree(), out_flat)
  
imports: from __future__ import annotations
  import functools
  import operator
  from typing import Callable
  from jax import lax
  from jax._src import api
  from jax._src import core
  from jax._src import custom_api_util
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import traceback_util
  from jax._src import tree_util
  from jax._src import util
  from jax._src.api_util import flatten_fun_nokwargs
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters.batching import not_mapped
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.interpreters import xla
  from jax._src.tree_util import tree_flatten, tree_map, tree_structure, tree_unflatten, treedef_tuple
label_desc: Use image node_id 3 for calling the custom_vmap obj's underlying member method code with example usage: obj.__call__() and returns: tree_unflatten