code: def linear_warmup(step, total_train_steps, steps_or_fraction=10000):
      """Linearly warm up the learning rate from 0.
  
      If steps_or_fraction > 1 , it is the absolute number of initial steps over
      which to warm up.  If it is <=1, then it is a fraction of the total number of
      training 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.
        steps_or_fraction: a number
  
  
      Returns:
        a tf.Scalar, the learning rate for the step.
      """
      warmup_steps = steps_or_fraction
      if steps_or_fraction <= 1:
          warmup_steps *= total_train_steps
      step = tf.cast(step, tf.float32)
      return tf.minimum(1.0, step / 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 3 for calling a global function with example usage: linear_warmup(step, total_train_steps, steps_or_fraction) and returns: tf