code: def _select_if_vmap_jvp(primals, tangents):
      pred, x, y = primals
      _, tx, ty = tangents
      assert x.shape == tx.aval.shape
      assert x.dtype == tx.aval.dtype
      assert y.shape == ty.aval.shape
      assert y.dtype == ty.aval.dtype
      out = _select_if_vmap(pred, x, y, makes_false_steps=False)
      if type(tx) is ad.Zero and type(ty) is ad.Zero:
          t_out = tx
      else:
          if type(tx) is ad.Zero:
              tx = jnp.zeros(
                  tx.aval.shape, tx.aval.dtype
              )  # pyright: ignore
          if type(ty) is ad.Zero:
              ty = jnp.zeros(
                  ty.aval.shape, ty.aval.dtype
              )  # pyright: ignore
          t_out = _select_if_vmap(pred, tx, ty, makes_false_steps=False)
      return out, t_out
  
imports: import itertools
  from typing import Any, TYPE_CHECKING, Union
  import jax
  import jax.core
  import jax.interpreters.ad
  import jax.interpreters.batching
  import jax.interpreters.mlir
  import jax.lax
  import jax.numpy
  import jax.tree_util
  from jaxtyping import Array, Bool
  from ..._filters import combine, is_array, partition
  from ..._module import field, Module
  from ..._pretty_print import tree_pformat
  from ..._tree import tree_at, tree_equal
  from ..._unvmap import unvmap_any
  from .._nontraceable import nonbatchable
  from .._primitive import create_vprim
label_desc: Use image node_id 3 for calling a global function with example usage: _select_if_vmap_jvp(primals, tangents) and returns: out,  t_out