added exisiting tests

Author: edawite

optax/schedules/_schedule_test.py

@@ -487,6 +487,7 @@ class WarmupCosineDecayTest(chex.TestCase):
   def test_monotonicity_and_exponent_ordering(self):
     init, peak, end = 0.0, 1.0, 0.1
     warmup_steps, decay_steps = 5, 25
+
     base = self.variant(_schedule.warmup_cosine_decay_schedule(
         init_value=init, peak_value=peak,
         warmup_steps=warmup_steps, decay_steps=decay_steps,
@@ -495,16 +496,82 @@ def test_monotonicity_and_exponent_ordering(self):
         init_value=init, peak_value=peak,
         warmup_steps=warmup_steps, decay_steps=decay_steps,
         end_value=end, exponent=2.0))
+
     steps = np.arange(decay_steps + 1)
-    vals, vals2 = base(steps), steep(steps)
-    for t in range(warmup_steps):
-      self.assertLess(float(vals[t]), float(vals[t + 1]))
-    self.assertAlmostEqual(float(vals[warmup_steps]), peak, places=6)
-    for t in range(warmup_steps, decay_steps):
-      self.assertGreaterEqual(float(vals[t]), float(vals[t + 1]))
-    self.assertAlmostEqual(float(vals[-1]), end, places=6)
-    for t in range(warmup_steps, decay_steps + 1):
-      self.assertLessEqual(float(vals2[t]), float(vals[t]) + 1e-12)
+    vals = base(steps)
+    vals2 = steep(steps)
+
+    with self.subTest("warmup increases"):
+      for t in range(warmup_steps):
+        self.assertLess(float(vals[t]), float(vals[t + 1]))
+
+    with self.subTest("peak at boundary"):
+      self.assertAlmostEqual(float(vals[warmup_steps]), peak, places=6)
+
+    with self.subTest("cosine decay nonincreasing and end value"):
+      for t in range(warmup_steps, decay_steps):
+        self.assertGreaterEqual(float(vals[t]), float(vals[t + 1]))
+      self.assertAlmostEqual(float(vals[-1]), end, places=6)
+
+    with self.subTest("exponent ordering (p=2 ≤ p=1)"):
+      for t in range(warmup_steps, decay_steps + 1):
+        self.assertLessEqual(float(vals2[t]), float(vals[t]) + 1e-12)
+
+  @chex.all_variants
+  def test_regression_values_subtests(self):
+
+    for name, init_value, peak_value, end_value in [
+        ("with end value", 10, 0.5, 1e-4),
+        ("without end value", 5, 3, 0.0),
+    ]:
+      with self.subTest(f"limits: {name}"):
+        schedule_fn = self.variant(_schedule.warmup_cosine_decay_schedule(
+            init_value=init_value,
+            peak_value=peak_value,
+            warmup_steps=100,
+            decay_steps=1000,
+            end_value=end_value,
+        ))
+        np.testing.assert_allclose(init_value, schedule_fn(0))
+        np.testing.assert_allclose(peak_value, schedule_fn(100))
+        np.testing.assert_allclose(end_value, schedule_fn(1000), rtol=1e-3)
+
+    with self.subTest("with exponent golden values"):
+      schedule_fn = self.variant(_schedule.warmup_cosine_decay_schedule(
+          init_value=0.2,
+          peak_value=1.21,
+          end_value=-3.0,
+          warmup_steps=50,
+          decay_steps=100,
+          exponent=2,
+      ))
+      output = schedule_fn(np.array([0, 10, 50, 75, 100]))
+      np.testing.assert_allclose(
+          output,
+          np.array([
+              0.20000004768371582,
+              0.4020000100135803,
+              1.2100000381469727,
+              -1.947500228881836,
+              -3.000000238418579,
+          ]),
+          rtol=1e-6, atol=1e-8,
+      )
+
+    with self.subTest("zero peak value"):
+      schedule_fn = self.variant(_schedule.warmup_cosine_decay_schedule(
+          init_value=0.2,
+          peak_value=0,
+          end_value=-3.0,
+          warmup_steps=50,
+          decay_steps=100,
+          exponent=2,
+      ))
+      output = schedule_fn(np.array([0, 10, 50, 75, 100]))
+      np.testing.assert_allclose(
+          output, np.array([0.2, 0.16, 0.0, 0.0, 0.0]),
+          rtol=1e-6, atol=1e-8,
+      )
 
   def test_raises_when_decay_equals_warmup(self):
     with self.assertRaises(ValueError):