class_docstr: ## Adam Optimizer for Half Precision Training
  
  We extend [Adam Optimizer](adam.html) but use FP32 to store gradients and moments.
class_docstr_tok: ["#","#","Adam","Optimizer","for","Half","Precision","Training","We","extend","[","Adam","Optimizer","]","(","adam.html",")","but","use","FP32","to","store","gradients","and","moments","."]
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}$
  
  All the state tensors use FP32.
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","}","$","All","the","state","tensors","use","FP32","."]
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}$
  
      All the state tensors use FP32.
      """
  
      # This is the number of optimizer steps taken on the parameter, $t$
      state["step"] = 0
      # Exponential moving average of gradients, $m_t$
      state["exp_avg"] = torch.zeros_like(
          param, memory_format=torch.preserve_format, dtype=torch.float
      )
      # Exponential moving average of squared gradient values, $v_t$
      state["exp_avg_sq"] = torch.zeros_like(
          param, memory_format=torch.preserve_format, dtype=torch.float
      )
      # Maintain a FP32 copy of the parameters
      state["fp32_copy"] = param.to(torch.float)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/optimizers/adam_fp16.py
imports: from typing import Dict, Tuple, Optional, Any
  import torch
  from torch import nn
  from torch.optim import Optimizer
  from torch.cuda.amp import grad_scaler
  from collections import defaultdict, abc
  from labml_nn.optimizers import WeightDecay
  from labml_nn.optimizers.adam import Adam
label_desc: Use image node_id 2 for calling the AdamFP16 obj's underlying member method code with example usage: obj.init_state(state, group, param) without return types