class_docstr: Variant of the Adam optimizer.
  
  It achieves fast convergence as Adam and generalization comparable to SGD.
  It adapts the step size depending on its "belief" in the gradient direction
  — the optimizer adaptively scales step size by the difference between the
  predicted and observed gradients.
  
  It implements the AdaBelief proposed by
  Juntang Zhuang et al. in [AdaBelief Optimizer: Adapting stepsizes by the
  belief in observed gradients](https://arxiv.org/abs/2010.07468).
  
  Example of usage:
  
  ```python
  opt = tfa.optimizers.AdaBelief(lr=1e-3)
  ```
  
  Note: `amsgrad` is not described in the original paper. Use it with
        caution.
  
  You can enable enable warmup by setting `total_steps` and
  `warmup_proportion`,
  and enable recitifcation as in RAdam by setting 'rectify':
  ```python
  opt = tfa.optimizers.AdaBelief(
      lr=1e-3,
      total_steps=10000,
      warmup_proportion=0.1,
      min_lr=1e-5,
      rectify=True,
  )
  ```
  
  In the above example, the learning rate will increase linearly
  from 0 to `lr` in 1000 steps, then decrease linearly from `lr` to `min_lr`
  in 9000 steps.
  
  Note 'rectify' is independent of 'warmup', you can choose any combinations.
  
  Lookahead, proposed by Michael R. Zhang et.al in the paper
  [Lookahead Optimizer: k steps forward, 1 step back]
  (https://arxiv.org/abs/1907.08610v1), can be integrated with AdaBelief,
  which is called 'ranger_adabelief' in the author's implementation
  https://github.com/juntang-zhuang/Adabelief-Optimizer.
  The mechanism can be enabled by using the lookahead wrapper. For example:
  
  ```python
  adabelief = tfa.optimizers.AdaBelief()
  ranger = tfa.optimizers.Lookahead(adabelief, sync_period=6, slow_step_size=0.5)
  ```
class_docstr_tok: ["Variant","of","the","Adam","optimizer",".","It","achieves","fast","convergence","as","Adam","and","generalization","comparable","to","SGD",".","It","adapts","the","step","size","depending","on","its","``","belief","''","in","the","gradient","direction","\u2014","the","optimizer","adaptively","scales","step","size","by","the","difference","between","the","predicted","and","observed","gradients",".","It","implements","the","AdaBelief","proposed","by","Juntang","Zhuang","et","al",".","in","[","AdaBelief","Optimizer",":","Adapting","stepsizes","by","the","belief","in","observed","gradients","]","(","https",":","\/\/arxiv.org\/abs\/2010.07468",")",".","Example","of","usage",":","``","`","python","opt","=","tfa.optimizers.AdaBelief","(","lr=1e-3",")","``","`","Note",":","`","amsgrad","`","is","not","described","in","the","original","paper",".","Use","it","with","caution",".","You","can","enable","enable","warmup","by","setting","`","total_steps","`","and","`","warmup_proportion","`",",","and","enable","recitifcation","as","in","RAdam","by","setting","'rectify","'",":","``","`","python","opt","=","tfa.optimizers.AdaBelief","(","lr=1e-3",",","total_steps=10000",",","warmup_proportion=0.1",",","min_lr=1e-5",",","rectify=True",",",")","``","`","In","the","above","example",",","the","learning","rate","will","increase","linearly","from","0","to","`","lr","`","in","1000","steps",",","then","decrease","linearly","from","`","lr","`","to","`","min_lr","`","in","9000","steps",".","Note","'rectify","'","is","independent","of","'warmup","'",",","you","can","choose","any","combinations",".","Lookahead",",","proposed","by","Michael","R.","Zhang","et.al","in","the","paper","[","Lookahead","Optimizer",":","k","steps","forward",",","1","step","back","]","(","https",":","\/\/arxiv.org\/abs\/1907.08610v1",")",",","can","be","integrated","with","AdaBelief",",","which","is","called","'ranger_adabelief","'","in","the","author","'s","implementation","https",":","\/\/github.com\/juntang-zhuang\/Adabelief-Optimizer",".","The","mechanism","can","be","enabled","by","using","the","lookahead","wrapper",".","For","example",":","``","`","python","adabelief","=","tfa.optimizers.AdaBelief","(",")","ranger","=","tfa.optimizers.Lookahead","(","adabelief",",","sync_period=6",",","slow_step_size=0.5",")","``","`"]
code: def _resource_apply_sparse(self, grad, var, indices):
      var_dtype = var.dtype.base_dtype
      lr_t = self._decayed_lr(var_dtype)
      wd_t = self._decayed_wd(var_dtype)
      beta_1_t = self._get_hyper("beta_1", var_dtype)
      beta_2_t = self._get_hyper("beta_2", var_dtype)
      epsilon_t = tf.convert_to_tensor(self.epsilon, var_dtype)
      local_step = tf.cast(self.iterations + 1, var_dtype)
      beta_1_power = tf.pow(beta_1_t, local_step)
      beta_2_power = tf.pow(beta_2_t, local_step)
  
      if self._initial_total_steps > 0:
          total_steps = self._get_hyper("total_steps", var_dtype)
          warmup_steps = total_steps * self._get_hyper(
              "warmup_proportion", var_dtype
          )
          min_lr = self._get_hyper("min_lr", var_dtype)
          decay_steps = tf.maximum(total_steps - warmup_steps, 1)
          decay_rate = (min_lr - lr_t) / decay_steps
          lr_t = tf.where(
              local_step <= warmup_steps,
              lr_t * (local_step / warmup_steps),
              lr_t
              + decay_rate
              * tf.minimum(local_step - warmup_steps, decay_steps),
          )
  
      sma_inf = 2.0 / (1.0 - beta_2_t) - 1.0
      sma_t = sma_inf - 2.0 * local_step * beta_2_power / (
          1.0 - beta_2_power
      )
  
      m = self.get_slot(var, "m")
      m_scaled_g_values = grad * (1 - beta_1_t)
      m_t = m.assign(m * beta_1_t, use_locking=self._use_locking)
      m_t = self._resource_scatter_add(m, indices, m_scaled_g_values)
      m_corr_t = m_t / (1.0 - beta_1_power)
  
      v = self.get_slot(var, "v")
      m_t_indices = tf.gather(m_t, indices)
      v_scaled_g_values = (
          tf.math.square(grad - m_t_indices) * (1 - beta_2_t)
          + epsilon_t
      )
      v_t = v.assign(v * beta_2_t, use_locking=self._use_locking)
      v_t = self._resource_scatter_add(v, indices, v_scaled_g_values)
  
      if self.amsgrad:
          vhat = self.get_slot(var, "vhat")
          vhat_t = vhat.assign(
              tf.maximum(vhat, v_t), use_locking=self._use_locking
          )
          v_corr_t = tf.sqrt(vhat_t / (1.0 - beta_2_power))
      else:
          vhat_t = None
          v_corr_t = tf.sqrt(v_t / (1.0 - beta_2_power))
  
      if self.rectify:
          r_t_numerator = (sma_t - 4.0) * (sma_t - 2.0) * sma_inf
          r_t_denominator = (sma_inf - 4.0) * (sma_inf - 2.0) * sma_t
          r_t = tf.sqrt(r_t_numerator / r_t_denominator)
          sma_threshold = self._get_hyper("sma_threshold", var_dtype)
          var_t = tf.where(
              sma_t >= sma_threshold,
              r_t * m_corr_t / (v_corr_t + epsilon_t),
              m_corr_t,
          )
      else:
          var_t = m_corr_t / (v_corr_t + epsilon_t)
  
      if self._has_weight_decay:
          var_t += wd_t * var
  
      var_update = self._resource_scatter_add(
          var, indices, tf.gather(-lr_t * var_t, indices)
      )
  
      updates = [var_update, m_t, v_t]
      if self.amsgrad:
          updates.append(vhat_t)
      return tf.group(*updates)
  
file_path: addons/tensorflow_addons/optimizers/adabelief.py
imports: import tensorflow
  from tensorflow_addons.utils.types import FloatTensorLike
  from tensorflow_addons.optimizers import KerasLegacyOptimizer
  from typing import Union, Callable, Dict
label_desc: Use image node_id 6 for calling the AdaBelief obj's underlying member method code with example usage: obj._resource_apply_sparse(grad, var, indices) and returns: tf