class_docstr: Applies a warmup schedule on a given learning rate decay schedule.
class_docstr_tok: ["Applies","a","warmup","schedule","on","a","given","learning","rate","decay","schedule","."]
code: def __call__(self, step: Union[int, tf.Tensor]) -> tf.Tensor:
      with tf.name_scope(self.name or "WarmUp") as name:
          # Implements linear 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 * warmup_percent_done
          )
          return tf.cond(
              global_step_float < warmup_steps_float,
              lambda: warmup_learning_rate,
              lambda: self.decay_schedule_fn(step),
              name=name,
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import tempfile
  from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
  import numpy
  import tensorflow
  from mediapipe.model_maker.python.core.data import dataset
  from mediapipe.model_maker.python.core.utils import quantization
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