code: def truncated_rsqrt(step, total_train_steps, warmup_steps=10000):
      """Noam's favorite learning-rate schedule.
  
      rsqrt(max(step_num, warmup_steps)
  
      Args:
        step: a tf.scalar representing the step we want the learning rate for.
        total_train_steps: a number, the total number of training steps.
        warmup_steps: a number
  
      Returns:
        a tf.Scalar, the learning rate for the step.
      """
      del total_train_steps
      step_num = tf.cast(step, tf.float32)
      return tf.math.rsqrt(tf.maximum(step_num, warmup_steps))
  
file_path: mesh/mesh_tensorflow/transformer/learning_rate_schedules.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import math
  import gin
  import tensorflow.compat.v1
label_desc: Use image node_id 4 for calling a global function with example usage: truncated_rsqrt(step, total_train_steps, warmup_steps) and returns: tf