class_docstr: Do various flavours of stochastic gradient descent, with first and
  second order momentum. Currently support 'vanilla' SGD, Adam, and RAdam.
class_docstr_tok: ["Do","various","flavours","of","stochastic","gradient","descent",",","with","first","and","second","order","momentum",".","Currently","support","'vanilla","'","SGD",",","Adam",",","and","RAdam","."]
code: def _radam(self, ops, weights, grad, lr_scale, key, nr_upd):
      if key not in self.mom1:
          self.mom1[key] = ops.alloc1f(weights.size)
      if key not in self.mom2:
          self.mom2[key] = ops.alloc1f(weights.size)
  
      weights_1D = ops.reshape1f(weights, weights.size)
      gradient_1D = ops.reshape1f(grad, grad.size)
  
      # While we port from the pytorch implementation, keep some of the same
      # naming
      state = {
          "step": self.nr_update[key],
          "exp_avg": self.mom1[key],
          "exp_avg_sq": self.mom2[key],
      }
      group = {
          "lr": self.learn_rate,
          "betas": [self.b1, self.b2],
          "eps": self.eps,
          "weight_decay": 0.0,
          "buffer": self._radam_buffer,
      }
      degenerated_to_sgd = True
  
      exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
      beta1, beta2 = group["betas"]
  
      # exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
      exp_avg_sq *= beta2
      exp_avg_sq += (1 - beta2) * (gradient_1D**2)
      # exp_avg.mul_(beta1).add_(1 - beta1, grad)
      exp_avg *= beta1
      exp_avg += (1 - beta1) * gradient_1D
  
      state["step"] += 1
      buffered = group["buffer"][int(state["step"] % 10)]
      if state["step"] == buffered[0]:
          N_sma, step_size = buffered[1], buffered[2]
      else:
          buffered[0] = state["step"]
          beta2_t = beta2 ** state["step"]
          N_sma_max = 2 / (1 - beta2) - 1
          N_sma = N_sma_max - 2 * state["step"] * beta2_t / (
              1 - beta2_t
          )
          buffered[1] = N_sma
  
          # more conservative since it's an approximated value
          if N_sma >= 5:
              step_size = math.sqrt(
                  (1 - beta2_t)
                  * (N_sma - 4)
                  / (N_sma_max - 4)
                  * (N_sma - 2)
                  / N_sma
                  * N_sma_max
                  / (N_sma_max - 2)
              ) / (1 - beta1 ** state["step"])
          elif degenerated_to_sgd:
              step_size = 1.0 / (1 - beta1 ** state["step"])
          else:
              step_size = -1
          buffered[2] = step_size
  
      # more conservative since it's an approximated value
      if N_sma >= 5:
          if group["weight_decay"] != 0:
              weights_1D += (
                  -group["weight_decay"] * group["lr"] * weights_1D
              )
          denom = ops.xp.sqrt(exp_avg_sq) + group["eps"]
          weights_1D += -step_size * group["lr"] * (exp_avg / denom)
      elif step_size > 0:
          if group["weight_decay"] != 0:
              weights_1D += (
                  -group["weight_decay"] * group["lr"] * weights_1D
              )
          weights_1D += -step_size * group["lr"] * exp_avg
      return (
          ops.reshape_f(weights_1D, weights.shape),
          ops.reshape_f(gradient_1D, grad.shape),
      )
  
file_path: thinc/thinc/optimizers.py
imports: import math
  from collections import defaultdict
  from typing import Dict, List, Optional, Tuple, Union, cast
  from .backends import get_array_ops
  from .config import registry
  from .types import FloatsXd, Generator
label_desc: Use image node_id 5 for calling the Optimizer obj's underlying member method code with example usage: obj._radam(ops, weights, grad, lr_scale, key, nr_upd) and returns: ops,  ops