code: def test_fn2fn_custom_idempotent():
      def fn(x):
          x = jnp.invert(x)
          x = eqxi.unvmap_any(x)
          x = jnp.invert(x)
          return x
  
      finalised_fn = eqxi.finalise_fn(fn)
      assert tree_allclose(fn(False), finalised_fn(False))
      assert tree_allclose(fn(True), finalised_fn(True))
  
      finalised_jaxpr = jax.make_jaxpr(finalised_fn)(True)
      finalised_jaxpr = cast(jax.core.ClosedJaxpr, finalised_jaxpr)
      finalised_finalised_jaxpr = jax.make_jaxpr(
          eqxi.finalise_fn(finalised_fn)
      )(True)
      finalised_finalised_jaxpr = cast(
          jax.core.ClosedJaxpr, finalised_finalised_jaxpr
      )
      _assert_jaxpr_equal(finalised_jaxpr, finalised_finalised_jaxpr)
      for eqn in finalised_jaxpr.eqns:
          assert eqn.primitive != eqxi.unvmap_any_p
  
      vmap_fn = jax.vmap(fn)
      finalised_vmap_fn = eqxi.finalise_fn(vmap_fn)
      for arg in (
          jnp.array([False, False]),
          jnp.array([False, True]),
          jnp.array([True, False]),
          jnp.array([True, True]),
      ):
          assert tree_allclose(vmap_fn(arg), finalised_vmap_fn(arg))
  
      finalised_vmap_jaxpr = jax.make_jaxpr(finalised_vmap_fn)(
          jnp.array([False, False])
      )
      finalised_vmap_jaxpr = cast(
          jax.core.ClosedJaxpr, finalised_vmap_jaxpr
      )
      finalised_finalised_vmap_jaxpr = jax.make_jaxpr(
          eqxi.finalise_fn(finalised_vmap_fn)
      )(jnp.array([False, False]))
      finalised_finalised_vmap_jaxpr = cast(
          jax.core.ClosedJaxpr, finalised_finalised_vmap_jaxpr
      )
      for eqn in finalised_vmap_jaxpr.eqns:
          assert eqn.primitive != eqxi.unvmap_any_p
      _assert_jaxpr_equal(
          finalised_vmap_jaxpr, finalised_finalised_vmap_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 7 for calling a global function with example usage: test_fn2fn_custom_idempotent() and returns: x