code: def _wrap_optimizer(opt: base.Optimizer) -> TransformedOptimizer:
      """Returns a functional optimizer."""
  
      def init_opt_fn(params):
          """Creates initial optimizer state."""
  
          def f(params):
              params = [p.deref() for p in sorted(params.keys())]
              updates = [tf.zeros_like(p) for p in params]
              for p, zero in zip(params, updates):
                  p.assign(zero)
              opt.apply(updates, params)
  
          f = haiku.transform_with_state(f)
  
          trainable, non_trainable = f.init(params)
          opt_state = _merge(
              {r: v for r, v in trainable.items() if r not in params},
              {
                  r: v
                  for r, v in non_trainable.items()
                  if r not in params
              },
          )
  
          return opt_state
  
      def apply_opt_fn(opt_state, updates, params):
          """Applies the optimizer and returns updated parameters and opt state."""
  
          def f(opt_state, params, updates):
              flat_params = [p.deref() for p in sorted(params)]
              updates = tree.flatten(updates)
              opt.apply(updates, flat_params)
              params = {r: r.deref().tensor_value for r in params}
              opt_state = {r: r.deref().tensor_value for r in opt_state}
              return params, opt_state
  
          f = haiku.transform_with_state(f)
  
          trainable_opt_state, non_trainable = _split_on_trainable(
              opt_state
          )
          trainable = _merge(params, trainable_opt_state)
          (params, opt_state), _ = f.apply(
              trainable, non_trainable, opt_state, params, updates
          )
          return params, opt_state
  
      return TransformedOptimizer(init=init_opt_fn, apply=apply_opt_fn)
  
file_path: sonnet/sonnet/src/functional/optimizers.py
imports: import collections
  import functools
  from typing import Callable, Type
  from sonnet.src import base
  from sonnet.src.functional import haiku
  import tensorflow
  import tree
label_desc: Use image node_id 4 for calling a global function with example usage: _wrap_optimizer(opt) and returns: TransformedOptimizer, opt_state, params,  opt_state, params,  opt_state