def __init__(self,
               learning_rate=0.01,
               momentum=0.9,
               second_moment_averaging=1.0,
               weight_decay=0.0,
               momentum_type=MomentumType.EMA):  # pylint: disable=useless-super-delegation
    """Create the SM3 optimizer.

    Memory-Efficient Adaptive Optimization.
    https://arxiv.org/abs/1901.11150

    Args:
      learning_rate: a postitive scalar value for the initial learning rate.
      momentum: optional, a positive scalar value for momentum
      second_moment_averaging: averaging of second moments (if 1.0, adds from
        begining of time like AdaGrad).
      weight_decay: Weight decay for regularizing the model.
      momentum_type: Nestrov, Heavy-Ball or EMA (Default).

    """
    self._has_momentum = momentum > 0.0
    self._momentum_type = momentum_type
    self._graft = second_moment_averaging != 1.0
    super().__init__(
        learning_rate=learning_rate,
        momentum=momentum,
        second_moment_averaging=second_moment_averaging,
        weight_decay=weight_decay,
    )
Create the SM3 optimizer.

    Memory-Efficient Adaptive Optimization.
    https://arxiv.org/abs/1901.11150

    Args:
      learning_rate: a postitive scalar value for the initial learning rate.
      momentum: optional, a positive scalar value for momentum
      second_moment_averaging: averaging of second moments (if 1.0, adds from
        begining of time like AdaGrad).
      weight_decay: Weight decay for regularizing the model.
      momentum_type: Nestrov, Heavy-Ball or EMA (Default).

    
trax/optimizers/sm3.py