Merge pull request #504 from brentyi:master

PiperOrigin-RevId: 513188521
google-deepmind · Mar 1, 2023 · 229b927 · 229b927
2 parents 4afbd1e + 263ed54
commit 229b927
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/optax/_src/schedule.py b/optax/_src/schedule.py
@@ -287,8 +287,8 @@ def piecewise_interpolate_schedule(
   else:
     boundaries, scales = (), ()
 
-  bounds = jnp.stack((0,) + boundaries)
-  values = jnp.cumprod(jnp.stack((init_value,) + scales))
+  bounds = np.stack((0,) + boundaries)
+  values = np.cumprod(np.stack((init_value,) + scales))
   interval_sizes = bounds[1:] - bounds[:-1]
 
   def schedule(count):