class_docstr: ## Sophia-G Optimizer
  
  We extend the class `GenericAdaptiveOptimizer` defined in [`__init__.py`](index.html)
  to implement the Sophia optimizer.
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}$
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}$
      """
  
      # 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
      )
      # Exponential moving average of Hessian diagonal, $h_t$
      state["hessian"] = torch.zeros_like(
          param, memory_format=torch.preserve_format
      )
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/optimizers/sophia.py
imports: from typing import Dict, Any, Tuple, Optional
  import torch
  from torch import nn
  from labml_nn.optimizers import GenericAdaptiveOptimizer, WeightDecay
label_desc: Use image node_id 2 for calling the Sophia obj's underlying member method code with example usage: obj.init_state(state, group, param) without return types