class_docstr: Variant of the Adam optimizer that handles sparse updates more
  efficiently.
  
  The original Adam algorithm maintains two moving-average accumulators for
  each trainable variable; the accumulators are updated at every step.
  This class provides lazier handling of gradient updates for sparse
  variables.  It only updates moving-average accumulators for sparse variable
  indices that appear in the current batch, rather than updating the
  accumulators for all indices. Compared with the original Adam optimizer,
  it can provide large improvements in model training throughput for some
  applications. However, it provides slightly different semantics than the
  original Adam algorithm, and may lead to different empirical results.
  
  Note, amsgrad is currently not supported and the argument can only be
  False.
class_docstr_tok: ["Variant","of","the","Adam","optimizer","that","handles","sparse","updates","more","efficiently",".","The","original","Adam","algorithm","maintains","two","moving-average","accumulators","for","each","trainable","variable",";","the","accumulators","are","updated","at","every","step",".","This","class","provides","lazier","handling","of","gradient","updates","for","sparse","variables",".","It","only","updates","moving-average","accumulators","for","sparse","variable","indices","that","appear","in","the","current","batch",",","rather","than","updating","the","accumulators","for","all","indices",".","Compared","with","the","original","Adam","optimizer",",","it","can","provide","large","improvements","in","model","training","throughput","for","some","applications",".","However",",","it","provides","slightly","different","semantics","than","the","original","Adam","algorithm",",","and","may","lead","to","different","empirical","results",".","Note",",","amsgrad","is","currently","not","supported","and","the","argument","can","only","be","False","."]
code: def _resource_apply_sparse(self, grad, var, indices):
      var_dtype = var.dtype.base_dtype
      lr_t = self._decayed_lr(var_dtype)
      beta_1_t = self._get_hyper("beta_1", var_dtype)
      beta_2_t = self._get_hyper("beta_2", var_dtype)
      local_step = tf.cast(self.iterations + 1, var_dtype)
      beta_1_power = tf.math.pow(beta_1_t, local_step)
      beta_2_power = tf.math.pow(beta_2_t, local_step)
      epsilon_t = tf.convert_to_tensor(self.epsilon, var_dtype)
      lr = lr_t * tf.math.sqrt(1 - beta_2_power) / (1 - beta_1_power)
  
      # \\(m := beta1 * m + (1 - beta1) * g_t\\)
      m = self.get_slot(var, "m")
      m_t_slice = (
          beta_1_t * tf.gather(m, indices) + (1 - beta_1_t) * grad
      )
      m_update_op = self._resource_scatter_update(m, indices, m_t_slice)
  
      # \\(v := beta2 * v + (1 - beta2) * (g_t * g_t)\\)
      v = self.get_slot(var, "v")
      v_t_slice = beta_2_t * tf.gather(v, indices) + (
          1 - beta_2_t
      ) * tf.math.square(grad)
      v_update_op = self._resource_scatter_update(v, indices, v_t_slice)
  
      # \\(variable += -learning_rate * m_t / (epsilon_t + sqrt(v_t))\\)
      var_slice = lr * m_t_slice / (tf.math.sqrt(v_t_slice) + epsilon_t)
      var_update_op = self._resource_scatter_sub(
          var, indices, var_slice
      )
  
      return tf.group(*[var_update_op, m_update_op, v_update_op])
  
file_path: addons/tensorflow_addons/optimizers/lazy_adam.py
imports: import importlib
  import tensorflow
  from tensorflow_addons.utils.types import FloatTensorLike
  from typeguard import typechecked
  from typing import Union, Callable
label_desc: Use image node_id 2 for calling the LazyAdam obj's underlying member method code with example usage: obj._resource_apply_sparse(grad, var, indices) and returns: tf