class_docstr: ## Rectified Adam Optimizer
  
  This class extends from AMSAdam optimizer defined in [`amsadam.py`](amsadam.html).
class_docstr_tok: ["#","#","Rectified","Adam","Optimizer","This","class","extends","from","AMSAdam","optimizer","defined","in","[","`","amsadam.py","`","]","(","amsadam.html",")","."]
docstr: ### Do the *RAdam* parameter update
  
  * `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}$
  * `m` and `v` are the uncorrected first and second moments $m_t$ and $v_t$;
    i.e. $\sigma(.)$ and $\psi(.)$ without bias correction
docstr_tok: ["#","#","#","Do","the","*","RAdam","*","parameter","update","*","`","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","}","$","*","`","m","`","and","`","v","`","are","the","uncorrected","first","and","second","moments","$","m_t","$","and","$","v_t","$",";","i.e",".","$","\\sigma","(",".",")","$","and","$","\\psi","(",".",")","$","without","bias","correction"]
code: def r_adam_update(
      self,
      state: Dict[str, any],
      group: Dict[str, any],
      param: torch.nn.Parameter,
      m: torch.Tensor,
      v: torch.Tensor,
  ):
      """
      ### Do the *RAdam* parameter update
  
      * `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}$
      * `m` and `v` are the uncorrected first and second moments $m_t$ and $v_t$;
        i.e. $\sigma(.)$ and $\psi(.)$ without bias correction
      """
  
      # Get $\beta_1$ and $\beta_2$
      beta1, beta2 = group["betas"]
      # Bias correction term for $\hat{m}_t$, $1 - \beta_1^t$
      bias_correction1 = 1 - beta1 ** state["step"]
      # Bias correction term for $\hat{v}_t$, $1 - \beta_2^t$
      bias_correction2 = 1 - beta2 ** state["step"]
  
      r = self.calc_rectification_term(beta2, state["step"])
  
      # Get learning rate
      lr = self.get_lr(state, group)
  
      # If $r_t$ is intractable
      if r is not None:
          # Whether to optimize the computation by combining scalar computations
          if self.optimized_update:
              # Denominator $\sqrt{v_t} + \hat{\epsilon}$
              denominator = v.sqrt().add_(group["eps"])
              # Step size $\alpha \sqrt{r_t} * \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t}$
              step_size = (
                  lr
                  * math.sqrt(bias_correction2)
                  * r
                  / bias_correction1
              )
              # Update parameters $\theta_t \leftarrow \theta_{t-1} - \alpha \sqrt{r_t} \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t} \cdot
              #  \frac{m_t}{\sqrt{v_t} + \hat{\epsilon}}$
              param.data.addcdiv_(m, denominator, value=-step_size)
          # Computation without optimization
          else:
              # Denominator  $\frac{\sqrt{v_t}}{\sqrt{1-\beta_2^t}} + \epsilon$
              denominator = (
                  v.sqrt() / math.sqrt(bias_correction2)
              ).add_(group["eps"])
              # Step size $\frac{\alpha \sqrt{r_t}}{1-\beta_1^t}$
              step_size = lr * r / bias_correction1
              # Update parameters $\theta_t \leftarrow \theta_{t-1} - \alpha \sqrt{r_t} \cdot
              # \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$
              param.data.addcdiv_(m, denominator, value=-step_size)
  
      # If $r_t$ is intractable do a SGD with momentum
      elif self.degenerated_to_sgd:
          # Step size $\frac{\alpha}{1-\beta_1^t}$
          step_size = lr / bias_correction1
          # Update parameters
          # $\theta_t \leftarrow \theta_{t-1} - \alpha \cdot \hat{m}_t$
          param.data.add_(m, alpha=-step_size)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/optimizers/radam.py
imports: import math
  from typing import Dict, Optional
  import torch
  from labml_nn.optimizers import WeightDecay
  from labml_nn.optimizers.amsgrad import AMSGrad
label_desc: Use image node_id 4 for calling the RAdam obj's underlying member method code with example usage: obj.r_adam_update(state, group, param, m, v) without return types