code: def get_learning_rate(learning_rate, hidden_size, learning_rate_warmup_steps):
    
    with tf.name_scope("learning_rate"):
      warmup_steps = tf.to_float(learning_rate_warmup_steps)
      step = tf.to_float(tf.train.get_or_create_global_step())
  
      learning_rate *= (hidden_size ** -0.5)
      # Apply linear warmup
      learning_rate *= tf.minimum(1.0, step / warmup_steps)
      # Apply rsqrt decay
      learning_rate *= tf.rsqrt(tf.maximum(step, warmup_steps))
  
      # Save learning rate value to TensorBoard summary.
      tf.summary.scalar("learning_rate", learning_rate)
  
      return learning_rate
docs: Calculate learning rate with linear warmup and rsqrt decay.