class_docstr: Adam optimizer; described in https://arxiv.org/abs/1412.6980.
  
  The update rule for time step :math:`t`, given gradients :math:`g_t` and
  "Stepsize" :math:`\alpha`, is:
  
  .. math::
      \hat{m}_t &\leftarrow \big(\beta_1 \cdot m_{t-1} + (1 - \beta_1) \cdot g_t\big)\ /\ (1 - \beta_1^t) \\
      \hat{v}_t &\leftarrow \big(\beta_2 \cdot m_{t-1} + (1 - \beta_2) \cdot g_t^2\big)\ /\ (1 - \beta_2^t) \\
      \theta_t  &\leftarrow \theta_{t-1} -\ \alpha \cdot \hat{m}_t / \big(\sqrt{\hat{v}_t} + \epsilon\big)
code: def update(self, step, grads, weights, slots, opt_params):
      m, v = slots
      learning_rate = opt_params["learning_rate"]
      weight_decay_rate = opt_params["weight_decay_rate"]
      b1 = opt_params["b1"]
      b2 = opt_params["b2"]
      eps = opt_params["eps"]
      m = (1 - b1) * grads + b1 * m  # First  moment estimate.
      v = (1 - b2) * (grads**2) + b2 * v  # Second moment estimate.
      mhat = m / (1 - b1 ** (step + 1))  # Bias correction.
      vhat = v / (1 - b2 ** (step + 1))
      new_weights = (
          (1 - weight_decay_rate) * weights
          - (learning_rate * mhat / (jnp.sqrt(vhat) + eps))
      ).astype(weights.dtype)
      return new_weights, (m, v)
  
imports: from trax.fastmath import numpy
  from trax.optimizers import base
label_desc: Use image node_id 3 for calling the Adam obj's underlying member method code with example usage: obj.update(step, grads, weights, slots, opt_params) and returns: new_weights,  m,  v