code: def warmup(n_warmup_steps, max_value):
      """Returns an LR schedule with linear warm-up followed by constant value.
  
      Args:
        n_warmup_steps: Number of steps during which the learning rate rises on
            a line connecting (0, 0) and (n_warmup_steps, max_value).
        max_value: Value for learning rate after warm-up has finished.
      """
      return _BodyAndTail(max_value, body_start=n_warmup_steps + 1)
  
file_path: trax/trax/supervised/lr_schedules.py
imports: import math
  import gin
  from trax.fastmath import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: warmup(n_warmup_steps, max_value) and returns: _BodyAndTail