code: def _select_if_vmap_impl(pred, x, y):
      # Not including the following, as it destroys performance on the GPU: it seems like
      # a copy of `x` is being made.
      #
      # msg = (
      #     "Internal error in Equinox. Please report a bug at "
      #     "https://github.com/patrick-kidger/equinox."
      # )
      # x = error_if(x, jnp.invert(pred), msg)
      return x
  
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 1 for calling a global function with example usage: _select_if_vmap_impl(pred, x, y) and returns: x