code: def test_warmup_and_rsqrt_decay(self):
      lr_fn = lr_schedules.warmup_and_rsqrt_decay(24, 0.25)
  
      # Warm-up.
      self.assertAlmostEqual(0.01, lr_fn(1))
      self.assertAlmostEqual(0.02, lr_fn(2))
      self.assertAlmostEqual(0.23, lr_fn(23))
      self.assertAlmostEqual(0.24, lr_fn(24))
  
      # Reciprocal square-root decay.
      self.assertAlmostEqual(0.25 * (5 / math.sqrt(25)), lr_fn(25))
      self.assertAlmostEqual(0.25 * (5 / math.sqrt(26)), lr_fn(26))
      self.assertAlmostEqual(0.25 * (5 / math.sqrt(27)), lr_fn(27))
      self.assertAlmostEqual(0.25 * (5 / math.sqrt(300)), lr_fn(300))
      self.assertAlmostEqual(0.25 * (5 / math.sqrt(4000)), lr_fn(4000))
      self.assertAlmostEqual(
          0.25 * (5 / math.sqrt(50000)), lr_fn(50000)
      )
  
file_path: trax/trax/supervised/lr_schedules_test.py
imports: import math
  from absl.testing import absltest
  from trax.supervised import lr_schedules
label_desc: Use image node_id 3 for calling the LRFunctionsTest obj's underlying member method code with example usage: obj.test_warmup_and_rsqrt_decay() without return types