code: def _addupdate_vmap(batched_args, batched_dims, *, tree):
      (axis_size,) = {
          x.shape[d]
          for x, d in zip(batched_args, batched_dims)
          if d is not batching.not_mapped
      }
      ref, val, *flat_idxs = batched_args
      ref_dim, val_dim, *flat_idx_dims = batched_dims
      indexers = tree_util.tree_unflatten(tree, flat_idxs)
      indexers_dims = tree_util.tree_unflatten(tree, flat_idx_dims)
  
      ref_is_batched = ref_dim is not batching.not_mapped
      val_is_batched = val_dim is not batching.not_mapped
      idx_is_batched = any(
          i_dim is not batching.not_mapped for i_dim in flat_idx_dims
      )
      if len(indexers) > 1:
          raise NotImplementedError(
              "Batching with multiple indexers not supported."
          )
      # TODO(sharadmv): handle vmap of multiple indexers
      indexers = tuple(
          _batch_indexer(
              indexer,
              dims,
              axis_size,
              ref.shape,
              ref_dim,
              idx_is_batched,
          )
          for indexer, dims in zip(indexers, indexers_dims)
      )
      flat_indexers, tree = tree_util.tree_flatten(indexers)
      if (ref_is_batched or idx_is_batched) and not val_is_batched:
          val = batching.broadcast(val, axis_size, 0)
      if val_is_batched:
          val = batching.moveaxis(val, val_dim, 0)
      return addupdate_p.bind(ref, val, *flat_indexers, tree=tree), []
  
imports: from __future__ import annotations
  from functools import partial
  from typing import Any, Union
  import numpy
  from jax._src import ad_util
  from jax._src import core
  from jax._src import pretty_printer
  from jax._src import tree_util
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import partial_eval
  from jax._src.interpreters import mlir
  from jax._src.lax import lax
  from jax._src.typing import Array
  from jax._src.state import indexing
  from jax._src.state.types import AbstractRef, RefView, ReadEffect, WriteEffect, AccumEffect
  from jax._src.util import safe_map, safe_zip
label_desc: Use image node_id 32 for calling a global function with example usage: _addupdate_vmap(batched_args, batched_dims) and returns: addupdate_p,  list