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 update(self, step, grads, weights, velocity, opt_params):
      del step
      v = velocity
      mu = opt_params["mass"]
      alpha = opt_params["weight_decay_rate"]
      epsilon = opt_params["learning_rate"]
  
      new_v = mu * v + grads
      if self._nesterov:
          weight_update = mu * new_v + grads
      else:
          weight_update = new_v
      new_weights = (1 - alpha) * weights - epsilon * weight_update
  
      new_weights = new_weights.astype(weights.dtype)
      return (new_weights, new_v)
  
imports: from trax.fastmath import numpy
  from trax.optimizers import base
label_desc: Use image node_id 3 for calling the Momentum obj's underlying member method code with example usage: obj.update(step, grads, weights, velocity, opt_params) and returns: new_weights,  new_v