code: def finalise_fn(fn):
      """Wraps a function to perform finalisation. (In a manner similar to e.g. `jax.vmap`
      wrapping a function to perform batching.)
      """
  
      def _finalise_fn(*args):
          jaxpr, struct = jax.make_jaxpr(fn, return_shape=True)(
              *args
          )  # pyright: ignore
          flat_args = jtu.tree_leaves(args)
          out = finalise_eval_jaxpr(
              jaxpr.jaxpr, jaxpr.consts, *flat_args
          )
          treedef = jtu.tree_structure(struct)
          return jtu.tree_unflatten(treedef, out)
  
      return _finalise_fn
  
imports: import functools
  from collections.abc import Callable
  from typing import Any, cast, Literal, overload, Union
  import jax
  import jax.core
  import jax.custom_derivatives
  import jax.tree_util
  from jaxtyping import PyTree
  from .._unvmap import unvmap_all_p, unvmap_any_p, unvmap_max_p
  from ..debug._announce_transform import announce_jaxpr_p
  from ._loop import maybe_set_p, select_if_vmap_p
  from ._noinline import noinline_p
  from ._nontraceable import nonbatchable_p, nondifferentiable_backward_p, nontraceable_p
label_desc: Use image node_id 10 for calling a global function with example usage: finalise_fn(fn) and returns: _finalise_fn, jtu