Merge pull request #1045 from google-deepmind:fabianp-patch-14

PiperOrigin-RevId: 671760146

optax/schedules/_schedule.py

@@ -60,7 +60,7 @@ def polynomial_schedule(
     transition_steps: int,
     transition_begin: int = 0
 ) -> base.Schedule:
-  """Constructs a schedule with polynomial transition from init to end value.
+  r"""Constructs a schedule with polynomial transition from init to end value.
 
   Examples:
     >>> schedule_fn = optax.polynomial_schedule(
@@ -70,6 +70,27 @@ def polynomial_schedule(
     >>> schedule_fn(100)  # learning rate on the last iteration
     Array(0.01, dtype=float32, weak_type=True)
 
+    The following example uses a non-zero `transition_begin`. In this case
+    the learning rate is kept constant for the first `transition_begin`
+    iterations:
+    >>> schedule_fn = optax.polynomial_schedule(
+    ...    init_value=1.0,
+    ...    end_value=0.01,
+    ...    transition_steps=100,
+    ...    transition_begin=5,
+    ...    power=2,
+    ... )
+    >>> counts = [0, 5, 6, 104, 105, 110]
+    >>> print(
+    ...    *[f'count:{i} value:{schedule_fn(i):.4f}' for i in counts],
+    ...    sep='\n')
+    count:0 value:1.0000
+    count:5 value:1.0000
+    count:6 value:0.9803
+    count:104 value:0.0101
+    count:105 value:0.0100
+    count:110 value:0.0100
+
   Args:
     init_value: initial value for the scalar to be annealed.
     end_value: end value of the scalar to be annealed.