code: def finalise_make_jaxpr(fn, *, return_shape: bool = False):
      """As `jax.make_jaxpr`, but finalises the result."""
  
      def _finalise_make_jaxpr(*args):
          jaxpr_struct = jax.make_jaxpr(
              fn, return_shape=return_shape
          )(  # pyright: ignore
              *args
          )
          if return_shape:
              jaxpr_struct = cast(
                  tuple[jax.core.ClosedJaxpr, Any], jaxpr_struct
              )
              jaxpr, struct = jaxpr_struct
              jaxpr = finalise_jaxpr(jaxpr)
              return jaxpr, struct
          else:
              jaxpr_struct = cast(jax.core.ClosedJaxpr, jaxpr_struct)
              jaxpr = finalise_jaxpr(jaxpr_struct)
              return jaxpr
  
      return _finalise_make_jaxpr
  
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 14 for calling a global function with example usage: finalise_make_jaxpr(fn) and returns: _finalise_make_jaxpr, jaxpr,  struct, jaxpr