text: import collections
  import typing
  
  from . import base
  from . import schedulers
  
  
  __all__ = ['AMSGrad']
  
  
  class AMSGrad(base.Optimizer):
      """AMSGrad optimizer.
  
      Parameters:
          lr: The learning rate.
          beta_1
          beta_2
          eps
          correct_bias
  
      Attributes:
          m (collections.defaultdict)
          v (collections.defaultdict)
          v_hat (collections.defaultdict)
  
      Example:
  
          >>> from creme import datasets
          >>> from creme import evaluate
          >>> from creme import linear_model
          >>> from creme import metrics
          >>> from creme import optim
          >>> from creme import preprocessing
  
          >>> dataset = datasets.Phishing()
          >>> optimizer = optim.AMSGrad()
          >>> model = (
          ...     preprocessing.StandardScaler() |
          ...     linear_model.LogisticRegression(optimizer)
          ... )
          >>> metric = metrics.F1()
  
          >>> evaluate.progressive_val_score(dataset, model, metric)
          F1: 0.865724
  
      References:
          1. [Reddi, S.J., Kale, S. and Kumar, S., 2019. On the convergence of adam and beyond. arXiv preprint arXiv:1904.09237](https://arxiv.org/pdf/1904.09237.pdf)
  
      """
  
      def __init__(self, lr: typing.Union[float, schedulers.Scheduler] = .1,
                   beta_1=.9, beta_2=.999, eps=1e-8, correct_bias=True):
          super().__init__(lr)
          self.beta_1 = beta_1
          self.beta_2 = beta_2
          self.eps = eps
          self.correct_bias = correct_bias
          self.m = collections.defaultdict(float)
          self.v = collections.defaultdict(float)
          self.v_hat = collections.defaultdict(float)
  
      def _update_after_pred(self, w, g):
  
          lr = self.learning_rate
  
          if self.correct_bias:
              # Correct bias for `v`
              lr *= (1 - self.beta_2 ** (self.n_iterations + 1)) ** .5
              # Correct bias for `m`
              lr /= 1 - self.beta_1 ** (self.n_iterations + 1)
  
          for i, gi in g.items():
              self.m[i] = self.beta_1 * self.m[i] + (1 - self.beta_1) * gi
              self.v[i] = self.beta_2 * self.v[i] + (1 - self.beta_2) * gi ** 2
              self.v_hat[i] = max(self.v_hat[i], self.v[i])
  
              w[i] -= lr * self.m[i] / (self.v_hat[i] ** 0.5 + self.eps)
  
          return w
  