code: def optimizer(
      cls: Type[base.Optimizer],
  ) -> Callable[..., TransformedOptimizer]:
      """Converts a snt.Optimizer subclass into a functional optimizer.
  
      To wrap a Sonnet optimizer class simply pass it to :func:`optimizer`:
  
      >>> adam = snt.functional.optimizer(snt.optimizers.Adam)
  
      This will give you back a function that drives the constructor of the
      optimizer and returns a pair of functions that give you the optimizer state
      and a way to apply it:
  
      >>> optimizer = adam(learning_rate=0.01)
  
      NOTE: We provide convenience wrappers for the builtin optimizers so you can
      just use `opt = snt.functional.adam(learning_rate=0.01)` if you prefer:
  
      >>> optimizer = snt.functional.adam(learning_rate=0.01)
  
      To make this example useful lets create a simple network to test:
  
      >>> with snt.functional.variables():
      ...   net = snt.nets.MLP([100, 10])
  
      >>> def loss_fn(images, labels):
      ...   logits = net(images)
      ...   x_ent = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
      ...                                                          labels=labels)
      ...   loss = tf.reduce_mean(x_ent)
      ...   return loss
  
      >>> loss_fn = snt.functional.transform(loss_fn)
  
      >>> x = tf.ones([1, 1])
      >>> y = tf.constant([1])
      >>> params = loss_fn.init(x, y)
  
      To get the initial state of our optimizer (e.g. m/v terms in Adam) we need to
      run the `optimizer.init` function:
  
      >>> opt_state = optimizer.init(params)
  
      Now we can run a single training step by taking gradients of our network and
      applying one step of our optimizer:
  
      >>> grad_apply_net = snt.functional.grad(loss_fn.apply)
  
      >>> def train_step(x, y, params, opt_state):
      ...   grads = grad_apply_net(params, x, y)
      ...   params, opt_state = optimizer.apply(opt_state, grads, params)
      ...   return params, opt_state
  
      Teach the network to always predict one:
  
      >>> target = tf.constant([1])
      >>> dataset = [(tf.random.normal([1, 1]), target) for _ in range(10)]
      >>> for x, y in dataset:
      ...   params, opt_state = train_step(x, y, params, opt_state)
  
      Args:
        cls: A :class:`~sonnet.Optimizer` subclass to functionalize.
  
      Returns:
        A transformed optimizer with `init` and `apply`. See docstring for details.
      """
  
      @functools.wraps(cls.__init__)
      def wrapper(*args, **kwargs):
          with haiku.variables():
              opt = cls(
                  *args, **kwargs
              )  # pytype: disable=not-instantiable
          return _wrap_optimizer(opt)
  
      return wrapper
  
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 1 for calling a global function with example usage: optimizer(cls) and returns: wrapper, _wrap_optimizer