class_docstr: Applies a warmup schedule on a given learning rate decay schedule.
  
  Args:
      initial_learning_rate (`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 (`Callable`):
          The schedule function to apply after the warmup for the rest of training.
      warmup_steps (`int`):
          The number of steps for the warmup part of training.
      power (`float`, *optional*, defaults to 1):
          The power to use for the polynomial warmup (defaults is a linear warmup).
      name (`str`, *optional*):
          Optional name prefix for the returned tensors during the schedule.
class_docstr_tok: ["Applies","a","warmup","schedule","on","a","given","learning","rate","decay","schedule",".","Args",":","initial_learning_rate","(","`","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","(","`","Callable","`",")",":","The","schedule","function","to","apply","after","the","warmup","for","the","rest","of","training",".","warmup_steps","(","`","int","`",")",":","The","number","of","steps","for","the","warmup","part","of","training",".","power","(","`","float","`",",","*","optional","*",",","defaults","to","1",")",":","The","power","to","use","for","the","polynomial","warmup","(","defaults","is","a","linear","warmup",")",".","name","(","`","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
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