code: def transform_with_state(f) -> TransformedWithState:
      r"""Like :func:`transform` but supporting non-trainable state.
  
      See :func:`transform` for more details.
  
      It is possible for the network to maintain internal state (e.g. for a module
      like `BatchNorm` that may want to maintain a moving average):
  
      >>> with snt.functional.variables():
      ...   ema = snt.ExponentialMovingAverage(decay=0.5)
  
      >>> f = snt.functional.transform_with_state(ema)
  
      When initializing this network we are returned the parameters (any "trainable"
      :tf:`Variable`\ s) and all other state (any non-trainable :tf:`Variable`\ s):
  
      >>> params, state = f.init(3.0)
      >>> params
      {}
      >>> state
      {<...>: <tf.Tensor: ... numpy=0>,
       <...>: <tf.Tensor: ... numpy=0.0>,
       <...>: <tf.Tensor: ... numpy=0.0>}
  
      To apply the network we simply call it and get back updated values for our
      non-trainable state:
  
      >>> y, state = f.apply(params, state, 3.0)
      >>> y.numpy()
      3.0
  
      >>> y, state = f.apply(params, state, 6.0)
      >>> y.numpy()
      5.0
  
      Args:
        f: A function closing over `Module` instances.
  
      Returns:
        A transformed function with `init` and `apply`. See docstring for details.
      """
  
      def init_fn(*args, **kwargs):
          """Applies `f(*a, **k)` and extracts initial variable values."""
          with create_tensor_variables(), track_new_variables() as new_variables, track_initial_state() as prev_var_state, track_tensor_variables() as tensor_variables:
              # NOTE: Intentionally discarding result.
              f(*args, **kwargs)
  
          params = initial_value_by_ref(
              v for v in tensor_variables if v.trainable
          )
          state = initial_value_by_ref(
              v for v in tensor_variables if not v.trainable
          )
  
          # Reset variable values.
          new_variables = {v.ref() for v in new_variables}
          for v in tensor_variables:
              r = v.ref()
              if r in new_variables:
                  # Variables created inside the function have their values nullified.
                  initial_tensor_value, tensor_value = None, None
              else:
                  # Variables that already existed have their value reset.
                  initial_tensor_value, tensor_value = prev_var_state[r]
              v.initial_tensor_value = initial_tensor_value
              v.tensor_value = tensor_value
  
          return params, state
  
      def apply_fn(params, state, *args, **kwargs):
          """Applies `f(*a, **k)` with variable values passed in."""
          initial_values = {}
          for r, t in itertools.chain(params.items(), state.items()):
              v = r.deref()
              initial_values[r] = (
                  v.tensor_value,
                  v.initial_tensor_value,
              )
              v.assign(t)
  
          try:
              with track_new_variables() as new_variables:
                  out = f(*args, **kwargs)
              if new_variables:
                  raise ValueError(
                      "Apply function cannot create new variables."
                  )
              state = final_value_by_ref(
                  p.deref() for p in state.keys()
              )
              return out, state
  
          finally:
              # Reset values to their initial state.
              for r, (
                  tensor_value,
                  initial_tensor_value,
              ) in initial_values.items():
                  v = r.deref()
                  v.tensor_value = tensor_value
                  v.initial_tensor_value = initial_tensor_value
  
      return TransformedWithState(init=init_fn, apply=apply_fn)
  
imports: import collections
  import functools
  import itertools
  import threading
  import contextlib
  from sonnet.src.functional import utils
  import tensorflow
label_desc: Use image node_id 16 for calling a global function with example usage: transform_with_state(f) and returns: TransformedWithState, params,  state, out,  state