class_docstr: SM3 optimizer, as described in https://arxiv.org/abs/1901.11150.
class_docstr_tok: ["SM3","optimizer",",","as","described","in","https",":","\/\/arxiv.org\/abs\/1901.11150","."]
docstr: Handle various types of momentum.
docstr_tok: ["Handle","various","types","of","momentum","."]
code: def _momentum_update(self, g, m, beta1):
      """Handle various types of momentum."""
      if self._momentum_type == MomentumType.EMA:
          m = (1 - beta1) * g + beta1 * m
          update = m
      elif self._momentum_type == MomentumType.HEAVY_BALL:
          m = g + beta1 * m
          update = m
      elif self._momentum_type == MomentumType.NESTEROV:
          m = g + beta1 * m
          nesterov_m = g + beta1 * m
          update = nesterov_m
      else:
          assert False, "Unknown momentum_type."
      return m, update
  
file_path: trax/trax/optimizers/sm3.py
imports: import enum
  from trax.fastmath import numpy
  from trax.optimizers import base
label_desc: Use image node_id 3 for calling the SM3 obj's underlying member method code with example usage: obj._momentum_update(g, m, beta1) and returns: m,  update