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 _adam(self, ops, weights, gradient, lr_scale, key, nr_upd):
      weights_1D = ops.reshape1f(weights, weights.size)
      gradient_1D = ops.reshape1f(gradient, gradient.size)
      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)
      mom1 = self.mom1[key]
      mom2 = self.mom2[key]
      b1 = self.b1
      b2 = self.b2
      fix1 = 1.0 - (b1**nr_upd)
      fix2 = 1.0 - (b2**nr_upd)
      lr = self.learn_rate * fix2**0.5 / fix1
      eps = self.eps
      # needs to be 1D going into the adam function
      weights_1D, gradient_1D, mom1, mom2 = ops.adam(
          weights_1D,
          gradient_1D,
          mom1,
          mom2,
          b1,
          b2,
          eps,
          lr * lr_scale,
      )
      self.mom1[key] = mom1
      self.mom2[key] = mom2
      return (
          ops.reshape_f(weights_1D, weights.shape),
          ops.reshape_f(gradient_1D, gradient.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 6 for calling the Optimizer obj's underlying member method code with example usage: obj._adam(ops, weights, gradient, lr_scale, key, nr_upd) and returns: ops,  ops