fix bugs in MultiStepLR (#11)

* fix bugs in `MultiStepLR` * update license
brainpy · Apr 22, 2024 · 4e9d942 · 4e9d942
1 parent 5b970fe
commit 4e9d942
Show file tree

Hide file tree

Showing 32 changed files with 75 additions and 39 deletions.
diff --git a/braintools/__init__.py b/braintools/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/functional/__init__.py b/braintools/functional/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/functional/_activations.py b/braintools/functional/_activations.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/functional/_normalization.py b/braintools/functional/_normalization.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/functional/_spikes.py b/braintools/functional/_spikes.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/init/__init__.py b/braintools/init/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/init/_base.py b/braintools/init/_base.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/init/_generic.py b/braintools/init/_generic.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/init/_random_inits.py b/braintools/init/_random_inits.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/init/_regular_inits.py b/braintools/init/_regular_inits.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/input/__init__.py b/braintools/input/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/input/currents.py b/braintools/input/currents.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/input/currents_test.py b/braintools/input/currents_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/__init__.py b/braintools/metric/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/_correlation.py b/braintools/metric/_correlation.py
@@ -1,4 +1,4 @@
-# Copyright 2024 BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/_correlation_test.py b/braintools/metric/_correlation_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/_fenchel_young.py b/braintools/metric/_fenchel_young.py
@@ -1,4 +1,4 @@
-# Copyright 2024 BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.
@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Fenchel-Young losses."""
 
 from typing import Any, Protocol
 
@@ -33,15 +32,15 @@ def __call__(self, scores, *args, **kwargs: Any):
 def make_fenchel_young_loss(
     max_fun: MaxFun
 ):
-  """Creates a Fenchel-Young loss from a max function.
+  """Creates a 2024 BDP Ecosystem from a max function.
 
   WARNING: The resulting loss accepts an arbitrary number of leading dimensions
   with the fy_loss operating over the last dimension. The jaxopt version of this
   function would instead flatten any vector in a single big 1D vector.
 
   Examples:
     Given a max function, e.g., the log sum exp, you can construct a
-    Fenchel-Young loss easily as follows:
+    2024 BDP Ecosystem easily as follows:
 
     >>> from jax.scipy.special import logsumexp
     >>> fy_loss = make_fy_loss(max_fun=logsumexp)
@@ -51,10 +50,10 @@ def make_fenchel_young_loss(
     <https://arxiv.org/pdf/1901.02324.pdf>`_, 2020
 
   Args:
-    max_fun: the max function on which the Fenchel-Young loss is built.
+    max_fun: the max function on which the 2024 BDP Ecosystem is built.
 
   Returns:
-    A Fenchel-Young loss function with the same signature.
+    A 2024 BDP Ecosystem function with the same signature.
   """
 
   vdot_last_dim = jnp.vectorize(jnp.vdot, signature="(n),(n)->()")

diff --git a/braintools/metric/_fenchel_young_test.py b/braintools/metric/_fenchel_young_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024 BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/_firings.py b/braintools/metric/_firings.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/_firings_test.py b/braintools/metric/_firings_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/_lfp.py b/braintools/metric/_lfp.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/_ranking.py b/braintools/metric/_ranking.py
@@ -1,4 +1,4 @@
-# Copyright 2024 BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/_ranking_test.py b/braintools/metric/_ranking_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024 BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/_regression.py b/braintools/metric/_regression.py
@@ -1,4 +1,4 @@
-# Copyright 2024 BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/_regression_test.py b/braintools/metric/_regression_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024 BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/_smoothing.py b/braintools/metric/_smoothing.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/_smoothing_test.py b/braintools/metric/_smoothing_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/metric/_util.py b/braintools/metric/_util.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/optim/__init__.py b/braintools/optim/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.

diff --git a/braintools/optim/_lr_scheduler.py b/braintools/optim/_lr_scheduler.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.
@@ -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):

diff --git a/braintools/optim/_lr_scheduler_test.py b/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))
diff --git a/braintools/optim/_sgd_optimizer.py b/braintools/optim/_sgd_optimizer.py
@@ -1,4 +1,4 @@
-# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+# 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.