class_docstr: SGD with momentum.
code: def apply_grad(self, grad, var):
      if grad is None:
          tf.logging.warning(
              "Gradient is None for variable %s", var.name
          )
          return []
  
      updates = []
      v = mtf.get_variable(
          var.mesh,
          var.name + "_momentum_v",
          var.shape,
          dtype=var.dtype,
          initializer=tf.zeros_initializer(),
          trainable=False,
      )
  
      with tf.variable_scope(var.name + "/sgd_momentum"):
          updates.append(
              mtf.assign(v, grad * self.lr + v * self.momentum)
          )
          updates.append(mtf.assign_sub(var, v))
  
      return updates
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import re
  import gin
  from mesh_tensorflow import layers
  from mesh_tensorflow import ops_with_redefined_builtins
  import tensorflow.compat.v1
label_desc: Use image node_id 4 for calling the MomentumOptimizer obj's underlying member method code with example usage: obj.apply_grad(grad, var) and returns: updates, list