class_docstr: ## Sophia-G Optimizer
  
  We extend the class `GenericAdaptiveOptimizer` defined in [`__init__.py`](index.html)
  to implement the Sophia optimizer.
docstr: ### Take an update step for a given parameter tensor
  
  * `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}$
  * `param` is the parameter tensor $     heta_{t-1}$
  
  We do the following parameter update,
  
  egin{align}
          heta_{t + 1} &\leftarrow        heta_t - \eta \cdot \operatorname{clip} igg(rac{m_t}{h_t + \epsilon}, 
  ho igg)
  \end{align}
code: def step_param(
      self,
      state: Dict[str, any],
      group: Dict[str, any],
      grad: torch.Tensor,
      param: torch.nn.Parameter,
  ):
      """
      ### Take an update step for a given parameter tensor
  
      * `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}$
      * `param` is the parameter tensor $\theta_{t-1}$
  
      We do the following parameter update,
  
      \begin{align}
      \theta_{t + 1} &\leftarrow \theta_t - \eta \cdot \operatorname{clip} \bigg(\frac{m_t}{h_t + \epsilon}, \rho \bigg)
      \end{align}
      """
  
      # Calculate weight decay
      grad = self.weight_decay(param, grad, group)
  
      # Get $\beta_1$ and $\beta_2$
      beta1, beta2 = group["betas"]
      # Get $\rho$
      rho = group["rho"]
  
      # Get $m_{t-1}$ and $h_{t}$
      m, hessian = state["exp_avg"], state["hessian"]
  
      # In-place calculation of $m_t$
      # $$m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) \cdot g_t$$
      m.mul_(beta1).add_(grad, alpha=1 - beta1)
  
      # Increment $t$ the number of optimizer steps
      state["step"] += 1
  
      # Get maximum learning rate $\eta \rho$
      lr = group["lr"]
  
      # $\eta$
      eta = lr / rho
  
      # $$\operatorname{clip} \bigg(\frac{m_t}{h_t + \epsilon}, \rho \bigg)$$
      ratio = (m / (hessian + group["eps"])).clamp(-rho, rho)
  
      # $$\theta_{t + 1} \leftarrow \theta_t - \eta \cdot \operatorname{clip} \bigg(\frac{m_t}{h_t + \epsilon}, \rho \bigg)$$
      param.data.add_(ratio, alpha=-eta)
  
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 4 for calling the Sophia obj's underlying member method code with example usage: obj.step_param(state, group, grad, param) without return types