code: def _vprim_jvp(
      primals,
      tangents,
      *,
      prim,
      __axis_size,
      __axis_name,
      __batch_axes,
      params,
  ):
      assert len(primals) == len(__batch_axes)
      assert len(tangents) == len(__batch_axes)
      tangents = [
          _resolve_zeros_t(t, b) for t, b in zip(tangents, __batch_axes)
      ]
      batch_axes_t = [
          _resolve_zeros_b(t, b) for t, b in zip(tangents, __batch_axes)
      ]
      jvp = ft.partial(_vprim_jvp_registry[prim], **params)
      jvp = jax.vmap(
          jvp,
          in_axes=(__batch_axes, batch_axes_t),
          axis_size=__axis_size,
          axis_name=__axis_name,
      )
      return jvp(primals, tangents)
  
imports: import functools
  from typing import Any
  import jax
  import jax.core
  import jax.interpreters.ad
  import jax.interpreters.batching
  import jax.interpreters.mlir
  import jax.tree_util
  from jaxtyping import PyTree
  from .._filters import combine, is_array, is_array_like, partition
  from .._tree import tree_equal
label_desc: Use image node_id 21 for calling a global function with example usage: _vprim_jvp(primals, tangents) and returns: jvp