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: ### Initialize a parameter state
  
  * `state` is the optimizer state of the parameter (tensor)
  * `group` stores optimizer attributes of the parameter group
  * `param` is the parameter tensor $     heta_{t-1}$
docstr_tok: ["#","#","#","Initialize","a","parameter","state","*","`","state","`","is","the","optimizer","state","of","the","parameter","(","tensor",")","*","`","group","`","stores","optimizer","attributes","of","the","parameter","group","*","`","param","`","is","the","parameter","tensor","$","heta_","{","t-1","}","$"]
code: def init_state(
      self,
      state: Dict[str, any],
      group: Dict[str, any],
      param: nn.Parameter,
  ):
      """
      ### Initialize a parameter state
  
      * `state` is the optimizer state of the parameter (tensor)
      * `group` stores optimizer attributes of the parameter group
      * `param` is the parameter tensor $\theta_{t-1}$
      """
  
      # Call `init_state` of Adam optimizer which we are extending
      super().init_state(state, group, param)
  
      # If `amsgrad` flag is `True` for this parameter group, we maintain the maximum of
      # exponential moving average of squared gradient
      if group["amsgrad"]:
          state["max_exp_avg_sq"] = torch.zeros_like(
              param, memory_format=torch.preserve_format
          )
  
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 2 for calling the AMSGrad obj's underlying member method code with example usage: obj.init_state(state, group, param) without return types