class_docstr: ## AMSGrad Optimizer
  
  This class extends from Adam optimizer defined in [`adam.py`](adam.html).
  Adam optimizer is extending the class `GenericAdaptiveOptimizer`
  defined in [`__init__.py`](index.html).
class_docstr_tok: ["#","#","AMSGrad","Optimizer","This","class","extends","from","Adam","optimizer","defined","in","[","`","adam.py","`","]","(","adam.html",")",".","Adam","optimizer","is","extending","the","class","`","GenericAdaptiveOptimizer","`","defined","in","[","`","__init__.py","`","]","(","index.html",")","."]
docstr: ### Calculate $m_t$ and and $v_t$ or $\max(v_1, v_2, ..., v_{t-1}, v_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","and","$","v_t","$","or","$","\\max","(","v_1",",","v_2",",","...",",","v_","{","t-1","}",",","v_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_mv(
      self,
      state: Dict[str, any],
      group: Dict[str, any],
      grad: torch.Tensor,
  ):
      """
      ### Calculate $m_t$ and and $v_t$ or $\max(v_1, v_2, ..., v_{t-1}, v_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 $m_t$ and $v_t$ from *Adam*
      m, v = super().get_mv(state, group, grad)
  
      # If this parameter group is using `amsgrad`
      if group["amsgrad"]:
          # Get $\max(v_1, v_2, ..., v_{t-1})$.
          #
          # 🗒 The paper uses the notation $\hat{v}_t$ for this, which we don't use
          # that here because it confuses with the Adam's usage of the same notation
          # for bias corrected exponential moving average.
          v_max = state["max_exp_avg_sq"]
          # Calculate $\max(v_1, v_2, ..., v_{t-1}, v_t)$.
          #
          # 🤔 I feel you should be taking / maintaining the max of the bias corrected
          # second exponential average of squared gradient.
          # But this is how it's
          # [implemented in PyTorch also](https://github.com/pytorch/pytorch/blob/19f4c5110e8bcad5e7e75375194262fca0a6293a/torch/optim/functional.py#L90).
          # I guess it doesn't really matter since bias correction only increases the value
          # and it only makes an actual difference during the early few steps of the training.
          torch.maximum(v_max, v, out=v_max)
  
          return m, v_max
      else:
          # Fall back to *Adam* if the parameter group is not using `amsgrad`
          return m, v
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/optimizers/amsgrad.py
imports: from typing import Dict
  import torch
  from torch import nn
  from labml_nn.optimizers import WeightDecay
  from labml_nn.optimizers.adam import Adam
label_desc: Use image node_id 3 for calling the AMSGrad obj's underlying member method code with example usage: obj.get_mv(state, group, grad) and returns: m,  v_max, m,  v