class_docstr: Momentum optimizer.
docstr: Updates variables and other state based on momentum (or Nesterov) SGD.
  
  Args:
     lr: the learning rate.
     grads: the gradients to apply.
     momentum: optional, override the default momentum.
docstr_tok: ["Updates","variables","and","other","state","based","on","momentum","(","or","Nesterov",")","SGD",".","Args",":","lr",":","the","learning","rate",".","grads",":","the","gradients","to","apply",".","momentum",":","optional",",","override","the","default","momentum","."]
code: def __call__(
      self,
      lr: float,
      grads: List[jn.ndarray],
      momentum: Optional[float] = None,
  ):
      """Updates variables and other state based on momentum (or Nesterov) SGD.
  
      Args:
         lr: the learning rate.
         grads: the gradients to apply.
         momentum: optional, override the default momentum.
      """
      assert len(grads) == len(
          self.train_vars
      ), "Expecting as many gradients as trainable variables"
      if momentum is None:
          momentum = self.momentum
      if self.nesterov:
          for g, p, m in zip(grads, self.train_vars, self.m):
              m.value = g + momentum * m.value
              p.value -= lr * (g + momentum * m.value)
      else:
          for g, p, m in zip(grads, self.train_vars, self.m):
              m.value = g + momentum * m.value
              p.value -= lr * m.value
  
file_path: objax/objax/optimizer/momentum.py
imports: from typing import List, Optional
  from jax import numpy
  from objax.module import Module, ModuleList
  from objax.util import class_name
  from objax.variable import TrainRef, StateVar, TrainVar, VarCollection
label_desc: Use image node_id 2 for calling the Momentum obj's underlying member method code with example usage: obj.__call__(lr, grads, momentum) without return types