def __init__(self,
               learning_rate=0.05,
               factored=True,
               multiply_by_parameter_scale=True,
               do_clipping=True,
               do_momentum=False,
               momentum_in_bfloat16=False,
               beta1=0.0,
               decay_rate=0.8,
               clipping_threshold=1.0,
               weight_decay_rate=1e-5,
               weight_decay_n_steps=0,
               epsilon1=1e-16,
               epsilon2=1e-3):
    """Create the Adafactor optimizer.

    Adafactor is described in https://arxiv.org/abs/1804.04235.

    Args:
      learning_rate: float: trax-provided learning rate.
      factored: boolean: whether to use factored second-moment estimator for 2d
        variables.
      multiply_by_parameter_scale: boolean: if True, then scale provided
        learning_rate by parameter norm. if False, provided learning_rate is
        absolute step size.
      do_clipping: whether to clip gradients; if True, set clipping_theshold.
      do_momentum: whether to use momentum; if True, set beta1.
      momentum_in_bfloat16: if True, store momentum in bfloat16 to save memory.
      beta1: a float value between 0 and 1, enables momentum and uses extra
        memory if nonzero!  Off by default.
      decay_rate: float: controls second-moment exponential decay schedule.
      clipping_threshold: an optional float >= 1, if None no update clipping.
      weight_decay_rate: rate at which to decay weights.
      weight_decay_n_steps: for how many steps to decay weights (always if None)
      epsilon1: Regularization constant for squared gradient.
      epsilon2: Regularization constant for parameter scale.
    """
    # These 4 parameters are not configurable once the class is created.
    self._factored = factored
    self._multiply_by_parameter_scale = multiply_by_parameter_scale
    self._do_clipping = do_clipping
    self._do_momentum = do_momentum
    self._momentum_in_bfloat16 = momentum_in_bfloat16
    # Dynamically configurable parameters will be passed to the update function.
    super().__init__(
        learning_rate=learning_rate,
        beta1=beta1,
        decay_rate=decay_rate,
        clipping_threshold=clipping_threshold,
        weight_decay_rate=weight_decay_rate,
        weight_decay_n_steps=weight_decay_n_steps,
        epsilon1=epsilon1,
        epsilon2=epsilon2,
    )
Create the Adafactor optimizer.

    Adafactor is described in https://arxiv.org/abs/1804.04235.

    Args:
      learning_rate: float: trax-provided learning rate.
      factored: boolean: whether to use factored second-moment estimator for 2d
        variables.
      multiply_by_parameter_scale: boolean: if True, then scale provided
        learning_rate by parameter norm. if False, provided learning_rate is
        absolute step size.
      do_clipping: whether to clip gradients; if True, set clipping_theshold.
      do_momentum: whether to use momentum; if True, set beta1.
      momentum_in_bfloat16: if True, store momentum in bfloat16 to save memory.
      beta1: a float value between 0 and 1, enables momentum and uses extra
        memory if nonzero!  Off by default.
      decay_rate: float: controls second-moment exponential decay schedule.
      clipping_threshold: an optional float >= 1, if None no update clipping.
      weight_decay_rate: rate at which to decay weights.
      weight_decay_n_steps: for how many steps to decay weights (always if None)
      epsilon1: Regularization constant for squared gradient.
      epsilon2: Regularization constant for parameter scale.
    
trax/optimizers/adafactor.py