code: def _select_if_vmap(pred, x, y, makes_false_steps):
      """As `lax.select(pred, x, y)` if `pred` is vmap'd. Not-vmap'd `pred` are assumed to
      be `True`, so that in this case `x` is returned unconditionally.
      """
      if makes_false_steps:
          return lax.select(pred, x, y)
      else:
          pred = fixed_asarray(pred)
          assert pred.shape == ()
          assert pred.dtype == jnp.bool_
          x = fixed_asarray(x)
          y = fixed_asarray(y)
          assert x.shape == y.shape
          assert x.dtype == y.dtype
          return select_if_vmap_p.bind(pred, x, y)
  
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 6 for calling a global function with example usage: _select_if_vmap(pred, x, y, makes_false_steps) and returns: lax, select_if_vmap_p