class_docstr: Adafactor.
code: def apply_grad(self, grad, var):
      if grad is None:
          tf.logging.warning(
              "Gradient is None for variable %s", var.name
          )
          return []
      # create slots
      grad = mtf.to_float(grad)
      factored_dims = self._factored_dims(var.shape)
      if factored_dims:
          d0, d1 = factored_dims
          vr_shape = var.shape - d0
          vc_shape = var.shape - d1
          vr = mtf.get_variable(
              var.mesh,
              var.name + "_slot_vr",
              vr_shape,
              initializer=tf.zeros_initializer(),
              trainable=False,
          )
          vc = mtf.get_variable(
              var.mesh,
              var.name + "_slot_vc",
              vc_shape,
              initializer=tf.zeros_initializer(),
              trainable=False,
          )
      else:
          v = mtf.get_variable(
              var.mesh,
              var.name + "_slot_v",
              var.shape,
              initializer=tf.zeros_initializer(),
              trainable=False,
          )
      if self._beta1:
          m = mtf.get_variable(
              var.mesh,
              var.name + "_slot_m",
              var.shape,
              initializer=tf.zeros_initializer(),
              trainable=False,
          )
  
      with tf.variable_scope(var.name + "/adafactor"):
          grad_squared = mtf.square(grad) + self._epsilon1
          decay_rate = self._decay_rate
          old_val = mtf.to_float(var.value)
          if self._multiply_by_parameter_scale and not any(
              [
                  s in var.name
                  for s in self._exclude_from_parameter_scale
              ]
          ):
              update_scale = (
                  self._parameter_scale(old_val) * self._learning_rate
              )
          else:
              update_scale = self._learning_rate
          mixing_rate = 1.0 - decay_rate
          updates = []
          if factored_dims:
              grad_squared_row_mean = mtf.reduce_mean(
                  grad_squared, output_shape=vr_shape
              )
              grad_squared_col_mean = mtf.reduce_mean(
                  grad_squared, output_shape=vc_shape
              )
              new_vr = (
                  vr * decay_rate + grad_squared_row_mean * mixing_rate
              )
              new_vc = (
                  vc * decay_rate + grad_squared_col_mean * mixing_rate
              )
              vr_update = mtf.assign(vr, new_vr)
              vc_update = mtf.assign(vc, new_vc)
              updates.extend([vr_update, vc_update])
              long_term_mean = mtf.reduce_mean(new_vr, reduced_dim=d1)
              r_factor = mtf.rsqrt(new_vr / long_term_mean)
              c_factor = mtf.rsqrt(new_vc)
              x = grad * r_factor * c_factor
          else:
              new_v = v * decay_rate + grad_squared * mixing_rate
              v_update = mtf.assign(v, new_v)
              updates.append(v_update)
              x = grad * mtf.rsqrt(new_v)
          if self._clipping_threshold is not None:
              clipping_denom = mtf.maximum(
                  1.0, reduce_rms(x) / self._clipping_threshold
              )
              x /= clipping_denom
          subtrahend = x * update_scale
          if self._beta1:
              new_m = m * tf.constant(
                  self._beta1
              ) + subtrahend * tf.constant(1.0 - self._beta1)
              subtrahend = new_m
              updates.append(mtf.assign(m, new_m))
          # It is critical to use assign_sub instead of mtf.assign(var - subtrahend)
          #  for the case of bfloat16 activations, so as to avoid repeatedly
          #  rounding the slice value, which results in poor quality.
          var_update = mtf.assign_sub(var, subtrahend)
          updates.append(var_update)
          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 AdafactorOptimizer obj's underlying member method code with example usage: obj.apply_grad(grad, var) and returns: list, updates