code: def interactive_trace(
      state: fun_mc.State,
      fn: fun_mc.TransitionOperator,
      num_steps: fun_mc.IntArray,
      trace_mask: bool | fun_mc.BooleanNest = True,
      iteration_axis: int = 0,
      block_until_ready: bool = True,
      progress_bar_fn: Callable[[Iterable[Any]], Iterable[Any]] = (
          _tqdm_progress_bar_fn
      ),
  ) -> tuple[fun_mc.State, fun_mc.ArrayNest]:
      """Wrapped around fun_mc.trace, suited for interactive work.
  
      This is accomplished through unrolling fun_mc.trace, as well as optionally
      using a progress bar (TQDM by default).
  
      Args:
        state: A nest of `Array`s or None.
        fn: A `TransitionOperator`.
        num_steps: Number of steps to run the function for. Must be greater than 1.
        trace_mask: A potentially shallow nest with boolean leaves applied to the
          `extra` return value of `fn`. This controls whether or not to actually
          trace the quantities in `extra`. For subtrees of `extra` where the mask
          leaf is `True`, those subtrees are traced (i.e. the corresponding subtrees
          in `traces` will contain an extra leading dimension equalling
          `num_steps`). For subtrees of `extra` where the mask leaf is `False`,
          those subtrees are merely propagated, and their corresponding subtrees in
          `traces` correspond to their final value.
        iteration_axis: Integer. Indicates the axis of the trace outputs that should
          be flattened with the first axis. This is most useful when `fn` is
          `trace`. E.g. if the trace has shape `[num_steps, 2, 5]` and
          `iteration_axis=2`, the trace outputs will be reshaped/transposed to `[2,
          5 * num_steps]`. A value of 0 disables this operation.
        block_until_ready: Whether to wait for the computation to finish between
          steps. This results in smoother progress bars under, e.g., JAX.
        progress_bar_fn: A callable that will be called with an iterable with length
          `num_steps` and which returns another iterable with the same length. This
          will be advanced for every step taken. If None, no progress bar is
          shown. Default: `lambda it: tqdm.tqdm(it, leave=True)`.
  
      Returns:
        state: The final state returned by `fn`.
        traces: A nest with the same structure as the extra return value of `fn`,
          but with leaves replaced with stacked and unstacked values according to
          the `trace_mask`.
      """
      num_steps = int(num_steps)
  
      if progress_bar_fn is None:
          pbar = None
      else:
          pbar = iter(progress_bar_fn(range(num_steps)))
  
      def fn_with_progress(state):
          state, extra = fun_mc.call_transition_operator(fn, state)
          if block_until_ready:
              state, extra = util.block_until_ready((state, extra))
          if pbar is not None:
              try:
                  next(pbar)
              except StopIteration:
                  pass
          return [state], extra
  
      [state], trace = fun_mc.trace(
          # Wrap the state in a singleton list to simplify implementation of
          # `fn_with_progress`.
          state=[state],
          fn=fn_with_progress,
          num_steps=num_steps,
          trace_mask=trace_mask,
          unroll=True,
      )
  
      if iteration_axis != 0:
  
          def fix_part(x):
              x = util.move_axis(x, 0, iteration_axis - 1)
              x = jnp.reshape(
                  x,
                  tuple(x.shape[: iteration_axis - 1])
                  + (-1,)
                  + tuple(x.shape[iteration_axis + 1 :]),
              )
              return x
  
          trace = util.map_tree(fix_part, trace)
      return state, trace
  
imports: import functools
  import typing
  from typing import Any, Callable, Iterable, NamedTuple, Optional
  import numpy
  from fun_mc import backend
  from fun_mc import fun_mc_lib
  from fun_mc import malt
  from fun_mc import sga_hmc
  from fun_mc.malt import *
  from fun_mc.sga_hmc import *
label_desc: Use image node_id 8 for calling a global function with example usage: interactive_trace(state, fn, num_steps, trace_mask, iteration_axis, block_until_ready, progress_bar_fn) and returns: state,  trace, list,  extra, x