code: def warmup_and_rsqrt_decay(n_warmup_steps, max_value):
      """Returns an LR schedule with warm-up + reciprocal square root decay."""
      return _BodyAndTail(
          max_value, tail_start=n_warmup_steps + 1, tail_fn=_rsqrt
      )
  
file_path: trax/trax/supervised/lr_schedules.py
imports: import math
  import gin
  from trax.fastmath import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: warmup_and_rsqrt_decay(n_warmup_steps, max_value) and returns: _BodyAndTail