fix bugs in MultiStepLR

braintools/optim/_lr_scheduler.py

@@ -17,6 +17,7 @@
 
 from typing import Sequence, Union
 
+import numpy as np
 import braincore as bc
 import jax
 import jax.numpy as jnp
@@ -222,13 +223,13 @@ def __init__(
       'milestones should be a sequence of increasing integers.'
     )
     assert 1. >= gamma >= 0, 'gamma should be in the range [0, 1].'
-    self.milestones = jnp.asarray((-1,) + tuple(milestones), dtype=bc.environ.ditype())
+    self.milestones = jnp.asarray((-1,) + tuple(milestones) + (np.iinfo(np.int32).max,), dtype=bc.environ.ditype())
     self.gamma = gamma
 
   def __call__(self, i=None):
     i = (self.last_epoch.value + 1) if i is None else i
-    conditions = (i > self.milestones[:-1]) & i < self.milestones[1:]
-    p = jnp.where(conditions, jnp.arange(0, len(self.milestones)))
+    conditions = jnp.logical_and((i >= self.milestones[:-1]), (i < self.milestones[1:]))
+    p = jnp.argmax(conditions)
     return self.lr * self.gamma ** p
 
   def extra_repr(self):

braintools/optim/_lr_scheduler_test.py

@@ -0,0 +1,36 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+
+import unittest
+
+import jax.numpy as jnp
+
+import braintools as bt
+
+
+class TestMultiStepLR(unittest.TestCase):
+  def test1(self):
+    lr = bt.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1)
+    for i in range(40):
+      r = lr(i)
+      if i < 10:
+        self.assertEqual(r, 0.1)
+      elif i < 20:
+        self.assertTrue(jnp.allclose(r, 0.01))
+      elif i < 30:
+        self.assertTrue(jnp.allclose(r, 0.001))
+      else:
+        self.assertTrue(jnp.allclose(r, 0.0001))