def __init__(self, learning_rate=0.0001, weight_decay_rate=1e-5,  # pylint: disable=useless-super-delegation
               b1=0.9, b2=0.999, eps=1e-5, clip_grad_norm=None):
    r"""Creates an Adam optimizer.

    Args:
      learning_rate: Initial (unadapted) learning rate :math:`\alpha`; original
          paper calls this `Stepsize` and suggests .001 as a generally good
          value.
      weight_decay_rate: Fraction of prior weight values to subtract on each
          step; equivalent to multiplying each weight element by
          `1 - weight_decay_rate`. (This is not part of the core Adam
          algorithm.)
      b1: Exponential decay rate :math:`\beta_1` for first moment estimates.
      b2: Exponential decay rate :math:`\beta_2` for second moment estimates.
      eps: Small positive constant :math:`\epsilon` for numerical stability.
      clip_grad_norm: Threshold value above which gradient clipping occurs.
          (This is not part of the core Adam algorithm.)
    """
    super().__init__(
        learning_rate=learning_rate,
        weight_decay_rate=weight_decay_rate,
        b1=b1,
        b2=b2,
        eps=eps,
        clip_grad_norm=clip_grad_norm
    )
Creates an Adam optimizer.

    Args:
      learning_rate: Initial (unadapted) learning rate :math:`\alpha`; original
          paper calls this `Stepsize` and suggests .001 as a generally good
          value.
      weight_decay_rate: Fraction of prior weight values to subtract on each
          step; equivalent to multiplying each weight element by
          `1 - weight_decay_rate`. (This is not part of the core Adam
          algorithm.)
      b1: Exponential decay rate :math:`\beta_1` for first moment estimates.
      b2: Exponential decay rate :math:`\beta_2` for second moment estimates.
      eps: Small positive constant :math:`\epsilon` for numerical stability.
      clip_grad_norm: Threshold value above which gradient clipping occurs.
          (This is not part of the core Adam algorithm.)
    
trax/optimizers/adam.py