code: def test_warmup(self):
      lr_fn = lr_schedules.warmup(9, 0.01)
  
      # Linear warm-up.
      self.assertAlmostEqual(0.001, lr_fn(1))
      self.assertAlmostEqual(0.002, lr_fn(2))
      self.assertAlmostEqual(0.005, lr_fn(5))
      self.assertAlmostEqual(0.009, lr_fn(9))
  
      # Constant thereafter.
      self.assertAlmostEqual(0.01, lr_fn(10))
      self.assertAlmostEqual(0.01, lr_fn(11))
      self.assertAlmostEqual(0.01, lr_fn(20))
      self.assertAlmostEqual(0.01, lr_fn(300))
      self.assertAlmostEqual(0.01, lr_fn(4000))
  
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 1 for calling the LRFunctionsTest obj's underlying member method code with example usage: obj.test_warmup() without return types