code: def test_cond():
      def f(pred, x):
          return lax.cond(pred, eqxi.unvmap_all, eqxi.unvmap_any, x)
  
      single_jaxpr = eqxi.finalise_make_jaxpr(f)(True, True)
      batch_jaxpr = eqxi.finalise_make_jaxpr(f)(
          False, jnp.array([True, False])
      )
  
      for jaxpr in (single_jaxpr, batch_jaxpr):
          _assert_no_unvmap(jaxpr.jaxpr)
  
file_path: equinox/tests/test_finalise_jaxpr.py
imports: from typing import cast
  import equinox.internal
  import jax
  import jax.core
  import jax.lax
  import jax.numpy
  from .helpers import tree_allclose
label_desc: Use image node_id 9 for calling a global function with example usage: test_cond() and returns: lax