Merge pull request #1222 from dpoteryayev:patch-1

PiperOrigin-RevId: 734170565
google-deepmind · Mar 6, 2025 · 2e66ce8 · 2e66ce8
2 parents 9138512 + 1a3b3a2
commit 2e66ce8
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/optax/schedules/_schedule.py b/optax/schedules/_schedule.py
@@ -330,9 +330,8 @@ def cosine_decay_schedule(
 
   .. math::
     \begin{cases}
-      \frac{I (1 - \alpha)}{2}(1+\cos(\pi\,\frac{t}{T})^p) + I \alpha\, 
-      & \text{if } t \leq T \\
-      I \alpha, & \text{if } t > T 
+      \alpha I + (1 - \alpha) I  \left[ \frac{1}{2} \left( 1+ \cos \left( \pi\,\frac{t}{T} \right) \right) \right] ^p  & \text{, if } t \leq T \\
+      \alpha I  & \text{, if } t > T 
     \end{cases}
 
   where :math:`T` is the number of decay steps (``decay_steps``), :math:`p` is