class_docstr: ```
  Applies a warmup schedule on a given learning rate decay schedule.
  Args:
      initial_learning_rate (:obj:`float`):
          The initial learning rate for the schedule after the warmup (so this will be the learning rate at the end
          of the warmup).
      decay_schedule_fn (:obj:`Callable`):
          The schedule function to apply after the warmup for the rest of training.
      warmup_steps (:obj:`int`):
          The number of steps for the warmup part of training.
      power (:obj:`float`, `optional`, defaults to 1):
          The power to use for the polynomial warmup (defaults is a linear warmup).
      name (:obj:`str`, `optional`):
          Optional name prefix for the returned tensors during the schedule.
  ```
code: def __call__(self, step):
      with tf.name_scope(self.name or "WarmUp") as name:
          # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
          # learning rate will be `global_step/num_warmup_steps * init_lr`.
          global_step_float = tf.cast(step, tf.float32)
          warmup_steps_float = tf.cast(self.warmup_steps, tf.float32)
          warmup_percent_done = global_step_float / warmup_steps_float
          warmup_learning_rate = (
              self.initial_learning_rate
              * tf.math.pow(warmup_percent_done, self.power)
          )
          return tf.cond(
              global_step_float < warmup_steps_float,
              lambda: warmup_learning_rate,
              lambda: self.decay_schedule_fn(step - self.warmup_steps),
              name=name,
          )
  
imports: import re
  from typing import Callable, List, Optional, Union
  import tensorflow
  from packaging import version
label_desc: Use image node_id 2 for calling the WarmUp obj's underlying member method code with example usage: obj.__call__(step) and returns: tf