class_docstr: ## AdaBelief Optimizer
  
  This class extends from RAdam optimizer defined in [`radam.py`](radam.html).
class_docstr_tok: ["#","#","AdaBelief","Optimizer","This","class","extends","from","RAdam","optimizer","defined","in","[","`","radam.py","`","]","(","radam.html",")","."]
docstr: ### Calculate $m_t$ and $s_t$ or $\max(s_1, s_2, ..., s_{t-1}, s_t)$
  
  * `state` is the optimizer state of the parameter (tensor)
  * `group` stores optimizer attributes of the parameter group
  * `grad` is the current gradient tensor $g_t$ for the parameter $       heta_{t-1}$
docstr_tok: ["#","#","#","Calculate","$","m_t","$","and","$","s_t","$","or","$","\\max","(","s_1",",","s_2",",","...",",","s_","{","t-1","}",",","s_t",")","$","*","`","state","`","is","the","optimizer","state","of","the","parameter","(","tensor",")","*","`","group","`","stores","optimizer","attributes","of","the","parameter","group","*","`","grad","`","is","the","current","gradient","tensor","$","g_t","$","for","the","parameter","$","heta_","{","t-1","}","$"]
code: def get_ms(
      self,
      state: Dict[str, Any],
      group: Dict[str, Any],
      grad: torch.Tensor,
  ):
      """
      ### Calculate $m_t$ and $s_t$ or $\max(s_1, s_2, ..., s_{t-1}, s_t)$
  
      * `state` is the optimizer state of the parameter (tensor)
      * `group` stores optimizer attributes of the parameter group
      * `grad` is the current gradient tensor $g_t$ for the parameter $\theta_{t-1}$
      """
  
      # Get $\beta_1$ and $\beta_2$
      beta1, beta2 = group["betas"]
  
      # Get $m_{t-1}$ and $s_{t-1}$
      m, s = state["exp_avg"], state["exp_avg_var"]
  
      # In-place calculation of $m_t$
      # $$m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) \cdot g_t$$
      m.mul_(beta1).add_(grad, alpha=1 - beta1)
      # Difference between gradient and momentum
      grad_residual = grad - m
      # In-place calculation of $s_t$
      # $$s_t \leftarrow \beta_2 s_{t-1} + (1 - \beta_2) \cdot (g_t - m_t)^2$$
      s.mul_(beta2).addcmul_(
          grad_residual, grad_residual, value=1 - beta2
      )
  
      # If this parameter group is using `amsgrad`
      if group["amsgrad"]:
          # Get $\max(s_1, s_2, ..., s_{t-1})$.
          s_max = state["max_exp_avg_var"]
          # Calculate $\max(s_1, s_2, ..., s_{t-1}, s_t)$.
          torch.maximum(s_max, s, out=s_max)
  
          return m, s_max
      else:
          # $m_t$ and $s_t$ otherwise
          return m, s
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/optimizers/ada_belief.py
imports: from typing import Dict, Any
  import torch
  from torch import nn
  from labml_nn.optimizers import WeightDecay
  from labml_nn.optimizers.radam import RAdam
label_desc: Use image node_id 3 for calling the AdaBelief obj's underlying member method code with example usage: obj.get_ms(state, group, grad) and returns: m,  s_max, m,  s