Add warmup_constant_schedule to optax.__init__.py.

PiperOrigin-RevId: 691192970
google-deepmind · Oct 29, 2024 · 2badbe8 · 2badbe8
1 parent e22f71b
commit 2badbe8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/optax/__init__.py b/optax/__init__.py
@@ -207,6 +207,7 @@
 piecewise_interpolate_schedule = schedules.piecewise_interpolate_schedule
 polynomial_schedule = schedules.polynomial_schedule
 sgdr_schedule = schedules.sgdr_schedule
+warmup_constant_schedule = schedules.warmup_constant_schedule
 warmup_cosine_decay_schedule = schedules.warmup_cosine_decay_schedule
 warmup_exponential_decay_schedule = schedules.warmup_exponential_decay_schedule
 inject_stateful_hyperparams = schedules.inject_stateful_hyperparams