code: def stateful_branch(branch_fun):
      """Calls branch_fun passing internal state in and out."""
  
      @functools.wraps(branch_fun)
      def new_branch_fun(operand):
          state, operand = operand
          with temporary_internal_state(
              state
          ), base.push_jax_trace_level():
              out = branch_fun(*operand)
              reserve_up_to_full_rng_block()
              # TODO(tomhennigan) Return difference of state in/out here.
              return out, internal_state()
  
      return new_branch_fun
  
imports: import collections
  import collections.abc
  from collections.abc import Mapping, MutableMapping
  import functools
  import inspect
  from typing import Any, Callable, Optional, TypeVar
  from haiku._src import base
  import jax
  import jax.numpy
label_desc: Use image node_id 12 for calling a global function with example usage: stateful_branch(branch_fun) and returns: new_branch_fun, out,  internal_state