class_docstr: Adam optimizer.
docstr: Updates variables and other state based on Adam algorithm.
  
  Args:
      lr: the learning rate.
      grads: the gradients to apply.
      beta1: optional, override the default beta1.
      beta2: optional, override the default beta2.
code: def __call__(
      self,
      lr: float,
      grads: List[JaxArray],
      beta1: Optional[float] = None,
      beta2: Optional[float] = None,
  ):
      """Updates variables and other state based on Adam algorithm.
  
      Args:
          lr: the learning rate.
          grads: the gradients to apply.
          beta1: optional, override the default beta1.
          beta2: optional, override the default beta2.
      """
      assert len(grads) == len(
          self.train_vars
      ), "Expecting as many gradients as trainable variables"
      if beta1 is None:
          beta1 = self.beta1
      if beta2 is None:
          beta2 = self.beta2
      self.step.value += 1
      lr *= jn.sqrt(1 - beta2**self.step.value) / (
          1 - beta1**self.step.value
      )
      for g, p, m, v in zip(grads, self.train_vars, self.m, self.v):
          m.value += (1 - beta1) * (g - m.value)
          v.value += (1 - beta2) * (g**2 - v.value)
          p.value -= lr * m.value * functional.rsqrt(v.value + self.eps)
  
file_path: objax/objax/optimizer/adam.py
imports: from typing import List, Optional
  from jax import numpy
  from objax import functional
  from objax.module import Module, ModuleList
  from objax.typing import JaxArray
  from objax.util import class_name
  from objax.variable import TrainRef, StateVar, TrainVar, VarCollection
label_desc: Use image node_id 2 for calling the Adam obj's underlying member method code with example usage: obj.__call__(lr, grads, beta1, beta2) without return types