class_docstr: A momentum optimizer.
  
  This class implements two variants of momentum stochastic gradient descent
  (SGD): with and without the Nesterov correction. The implementation of the
  Nesterov update is based on the concepts in Sutskever et al. (2013)
  [http://jmlr.org/proceedings/papers/v28/sutskever13.pdf], reformulated in
  Bengio et al. (2012) [https://arxiv.org/abs/1212.0901], to work well with
  backpropagation (equations 6 and 7):
  
  .. math::
      v_t      &= \mu_{t-1}v_{t-1} - \epsilon_{t-1}\nabla f(\Theta_{t-1}) \\
      \Theta_t &= \Theta_{t-1} - \mu_{t-1} v_{t-1} + \mu_t v_t + v_t
  
  where :math:`\mu_{t-1}` is the momentum (decay) coefficient at time step
  :math:`t-1` and :math:`\epsilon_{t-1}` is the learning rate at :math:`t-1`.
  
  Note that the implementation below also includes a weight decay rate
  (:math:`\alpha`) on the parameters, independent of the Nesterov momentum.
code: def init(self, weights):
      return jnp.zeros_like(weights)
  
imports: from trax.fastmath import numpy
  from trax.optimizers import base
label_desc: Use image node_id 2 for calling the Momentum obj's underlying member method code with example usage: obj.init(weights) and returns: jnp