code: def test_checkpoint():
      @jax.checkpoint  # pyright: ignore
      def f(pred, x):
          pred = eqxi.unvmap_any(pred)
          return lax.cond(pred, lambda y: y, lambda y: y + 1, x)
  
      jaxpr = eqxi.finalise_make_jaxpr(f)(True, 1.0)
      _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 12 for calling a global function with example usage: test_checkpoint() and returns: lax