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: ### 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}$
      """
      state["step"] = 0
      # Exponential moving average of gradient values
      state["exp_avg"] = torch.zeros_like(
          param, memory_format=torch.preserve_format
      )
      # Exponential moving average of variance
      state["exp_avg_var"] = torch.zeros_like(
          param, memory_format=torch.preserve_format
      )
  
      # If `amsgrad` flag is `True` for this parameter group, we maintain the maximum of
      # exponential moving average of variance
      if group["amsgrad"]:
          # Maintains max of all exp. moving avg. of sq. grad. values
          state["max_exp_avg_var"] = torch.zeros_like(
              param, memory_format=torch.preserve_format
          )
  
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 2 for calling the AdaBelief obj's underlying member method code with example usage: obj.init_state(state, group, param) without return types