tensorflow_federated/python/learning/optimizers/adam.py [84:105]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    normalized_lr = lr * tf.math.sqrt(
        (1 - tf.math.pow(beta_2, step))) / (1 - tf.math.pow(beta_1, step))
    updated_weights = tf.nest.map_structure(
        lambda w, g, a, s: w - normalized_lr * a / (tf.math.sqrt(s) + epsilon),
        weights, gradients, updated_accumulator, updated_preconditioner)

    updated_state = collections.OrderedDict([
        (_LEARNING_RATE_KEY, lr),
        (_BETA_1_KEY, beta_1),
        (_BETA_2_KEY, beta_2),
        (_EPSILON_KEY, epsilon),
        (_STEP_KEY, step),
        (_ACCUMULATOR_KEY, updated_accumulator),
        (_PRECONDITIONER_KEY, updated_preconditioner),
    ])
    return updated_state, updated_weights


def _check_beta(beta):
  py_typecheck.check_type(beta, float)
  if beta < 0.0 or beta >= 1.0:
    raise ValueError('Beta must be equal to 0.0 or more, and less than 1.0.')
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



tensorflow_federated/python/learning/optimizers/yogi.py [92:113]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    normalized_lr = lr * tf.math.sqrt(
        (1 - tf.math.pow(beta_2, step))) / (1 - tf.math.pow(beta_1, step))
    updated_weights = tf.nest.map_structure(
        lambda w, g, a, s: w - normalized_lr * a / (tf.math.sqrt(s) + epsilon),
        weights, gradients, updated_accumulator, updated_preconditioner)

    updated_state = collections.OrderedDict([
        (_LEARNING_RATE_KEY, lr),
        (_BETA_1_KEY, beta_1),
        (_BETA_2_KEY, beta_2),
        (_EPSILON_KEY, epsilon),
        (_STEP_KEY, step),
        (_ACCUMULATOR_KEY, updated_accumulator),
        (_PRECONDITIONER_KEY, updated_preconditioner),
    ])
    return updated_state, updated_weights


def _check_beta(beta):
  py_typecheck.check_type(beta, float)
  if beta < 0.0 or beta >= 1.0:
    raise ValueError('Beta must be equal to 0.0 or more, and less than 1.0.')
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -