brainstate 0.1.0.post20250503__py2.py3-none-any.whl → 0.1.2__py2.py3-none-any.whl

brainstate/nn/_module_test.py

@@ -13,20 +13,17 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import unittest
 
 import jax.numpy as jnp
-import jaxlib.xla_extension
 
-import brainstate as bst
+import brainstate
 
 
 class TestDelay(unittest.TestCase):
     def test_delay1(self):
-        a = bst.State(bst.random.random(10, 20))
-        delay = bst.nn.Delay(a.value)
+        a = brainstate.State(brainstate.random.random(10, 20))
+        delay = brainstate.nn.Delay(a.value)
         delay.register_entry('a', 1.)
         delay.register_entry('b', 2.)
         delay.register_entry('c', None)
@@ -36,7 +33,7 @@ class TestDelay(unittest.TestCase):
             delay.register_entry('c', 10.)
 
     def test_rotation_delay(self):
-        rotation_delay = bst.nn.Delay(jnp.ones((1,)))
+        rotation_delay = brainstate.nn.Delay(jnp.ones((1,)))
         t0 = 0.
         t1, n1 = 1., 10
         t2, n2 = 2., 20
@@ -53,7 +50,7 @@ class TestDelay(unittest.TestCase):
         # print(rotation_delay.max_length)
 
         for i in range(100):
-            bst.environ.set(i=i)
+            brainstate.environ.set(i=i)
             rotation_delay.update(jnp.ones((1,)) * i)
             # print(i, rotation_delay.at('a'), rotation_delay.at('b'), rotation_delay.at('c2'), rotation_delay.at('c'))
             self.assertTrue(jnp.allclose(rotation_delay.at('a'), jnp.ones((1,)) * i))
@@ -61,7 +58,7 @@ class TestDelay(unittest.TestCase):
             self.assertTrue(jnp.allclose(rotation_delay.at('c'), jnp.maximum(jnp.ones((1,)) * i - n2, 0.)))
 
     def test_concat_delay(self):
-        rotation_delay = bst.nn.Delay(jnp.ones([1]), delay_method='concat')
+        rotation_delay = brainstate.nn.Delay(jnp.ones([1]), delay_method='concat')
         t0 = 0.
         t1, n1 = 1., 10
         t2, n2 = 2., 20
@@ -74,7 +71,7 @@ class TestDelay(unittest.TestCase):
 
         print()
         for i in range(100):
-            bst.environ.set(i=i)
+            brainstate.environ.set(i=i)
             rotation_delay.update(jnp.ones((1,)) * i)
             print(i, rotation_delay.at('a'), rotation_delay.at('b'), rotation_delay.at('c'))
             self.assertTrue(jnp.allclose(rotation_delay.at('a'), jnp.ones((1,)) * i))
@@ -83,40 +80,40 @@ class TestDelay(unittest.TestCase):
         # bst.util.clear_buffer_memory()
 
     def test_jit_erro(self):
-        rotation_delay = bst.nn.Delay(jnp.ones([1]), time=2., delay_method='concat', interp_method='round')
+        rotation_delay = brainstate.nn.Delay(jnp.ones([1]), time=2., delay_method='concat', interp_method='round')
         rotation_delay.init_state()
 
-        with bst.environ.context(i=0, t=0, jit_error_check=True):
+        with brainstate.environ.context(i=0, t=0, jit_error_check=True):
             rotation_delay.retrieve_at_time(-2.0)
-            with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+            with self.assertRaises(Exception):
                 rotation_delay.retrieve_at_time(-2.1)
             rotation_delay.retrieve_at_time(-2.01)
-            with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+            with self.assertRaises(Exception):
                 rotation_delay.retrieve_at_time(-2.09)
-            with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+            with self.assertRaises(Exception):
                 rotation_delay.retrieve_at_time(0.1)
-            with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+            with self.assertRaises(Exception):
                 rotation_delay.retrieve_at_time(0.01)
 
     def test_round_interp(self):
         for shape in [(1,), (1, 1), (1, 1, 1)]:
             for delay_method in ['rotation', 'concat']:
-                rotation_delay = bst.nn.Delay(jnp.ones(shape), time=2., delay_method=delay_method,
-                                              interp_method='round')
+                rotation_delay = brainstate.nn.Delay(jnp.ones(shape), time=2., delay_method=delay_method,
+                                                     interp_method='round')
                 t0, n1 = 0.01, 0
                 t1, n1 = 1.04, 10
                 t2, n2 = 1.06, 11
                 rotation_delay.init_state()
 
-                @bst.compile.jit
+                @brainstate.compile.jit
                 def retrieve(td, i):
-                    with bst.environ.context(i=i, t=i * bst.environ.get_dt()):
+                    with brainstate.environ.context(i=i, t=i * brainstate.environ.get_dt()):
                         return rotation_delay.retrieve_at_time(td)
 
                 print()
                 for i in range(100):
-                    t = i * bst.environ.get_dt()
-                    with bst.environ.context(i=i, t=t):
+                    t = i * brainstate.environ.get_dt()
+                    with brainstate.environ.context(i=i, t=t):
                         rotation_delay.update(jnp.ones(shape) * i)
                         print(i,
                               retrieve(t - t0, i),
@@ -131,22 +128,22 @@ class TestDelay(unittest.TestCase):
             for delay_method in ['rotation', 'concat']:
                 print(shape, delay_method)
 
-                rotation_delay = bst.nn.Delay(jnp.ones(shape), time=2., delay_method=delay_method,
-                                              interp_method='linear_interp')
+                rotation_delay = brainstate.nn.Delay(jnp.ones(shape), time=2., delay_method=delay_method,
+                                                     interp_method='linear_interp')
                 t0, n0 = 0.01, 0.1
                 t1, n1 = 1.04, 10.4
                 t2, n2 = 1.06, 10.6
                 rotation_delay.init_state()
 
-                @bst.compile.jit
+                @brainstate.compile.jit
                 def retrieve(td, i):
-                    with bst.environ.context(i=i, t=i * bst.environ.get_dt()):
+                    with brainstate.environ.context(i=i, t=i * brainstate.environ.get_dt()):
                         return rotation_delay.retrieve_at_time(td)
 
                 print()
                 for i in range(100):
-                    t = i * bst.environ.get_dt()
-                    with bst.environ.context(i=i, t=t):
+                    t = i * brainstate.environ.get_dt()
+                    with brainstate.environ.context(i=i, t=t):
                         rotation_delay.update(jnp.ones(shape) * i)
                         print(i,
                               retrieve(t - t0, i),
@@ -157,8 +154,8 @@ class TestDelay(unittest.TestCase):
                         self.assertTrue(jnp.allclose(retrieve(t - t2, i), jnp.maximum(jnp.ones(shape) * i - n2, 0.)))
 
     def test_rotation_and_concat_delay(self):
-        rotation_delay = bst.nn.Delay(jnp.ones((1,)))
-        concat_delay = bst.nn.Delay(jnp.ones([1]), delay_method='concat')
+        rotation_delay = brainstate.nn.Delay(jnp.ones((1,)))
+        concat_delay = brainstate.nn.Delay(jnp.ones([1]), delay_method='concat')
         t0 = 0.
         t1, n1 = 1., 10
         t2, n2 = 2., 20
@@ -175,7 +172,7 @@ class TestDelay(unittest.TestCase):
 
         print()
         for i in range(100):
-            bst.environ.set(i=i)
+            brainstate.environ.set(i=i)
             new = jnp.ones((1,)) * i
             rotation_delay.update(new)
             concat_delay.update(new)
@@ -186,17 +183,17 @@ class TestDelay(unittest.TestCase):
 
 class TestModule(unittest.TestCase):
     def test_states(self):
-        class A(bst.nn.Module):
+        class A(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.a = bst.State(bst.random.random(10, 20))
-                self.b = bst.State(bst.random.random(10, 20))
+                self.a = brainstate.State(brainstate.random.random(10, 20))
+                self.b = brainstate.State(brainstate.random.random(10, 20))
 
-        class B(bst.nn.Module):
+        class B(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
                 self.a = A()
-                self.b = bst.State(bst.random.random(10, 20))
+                self.b = brainstate.State(brainstate.random.random(10, 20))
 
         b = B()
         print()
@@ -207,5 +204,5 @@ class TestModule(unittest.TestCase):
 
 
 if __name__ == '__main__':
-    with bst.environ.context(dt=0.1):
+    with brainstate.environ.context(dt=0.1):
         unittest.main()

brainstate/nn/metrics.py

@@ -14,8 +14,6 @@
 # ==============================================================================
 
 
-from __future__ import annotations
-
 import typing as tp
 from dataclasses import dataclass
 from functools import partial

brainstate/optim/_base.py

@@ -14,8 +14,6 @@
 # ==============================================================================
 
 
-from __future__ import annotations
-
 from typing import Dict, Hashable
 
 from brainstate._state import State, StateDictManager

brainstate/optim/_lr_scheduler.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 
 # -*- coding: utf-8 -*-
-from __future__ import annotations
 
 from typing import Sequence, Union
 

brainstate/optim/_lr_scheduler_test.py

@@ -19,12 +19,12 @@ import unittest
 
 import jax.numpy as jnp
 
-import brainstate as bst
+import brainstate
 
 
 class TestMultiStepLR(unittest.TestCase):
     def test1(self):
-        lr = bst.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1)
+        lr = brainstate.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1)
         for i in range(40):
             r = lr(i)
             if i < 10:
@@ -37,7 +37,7 @@ class TestMultiStepLR(unittest.TestCase):
                 self.assertTrue(jnp.allclose(r, 0.0001))
 
     def test2(self):
-        lr = bst.compile.jit(bst.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1))
+        lr = brainstate.compile.jit(brainstate.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1))
         for i in range(40):
             r = lr(i)
             if i < 10:

brainstate/optim/_optax_optimizer.py

@@ -14,8 +14,6 @@
 # ==============================================================================
 
 
-from __future__ import annotations
-
 import importlib.util
 from typing import Hashable, Dict, Optional
 

brainstate/optim/_optax_optimizer_test.py

@@ -13,39 +13,38 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
 import jax
 import optax
 
-import brainstate as bst
+import brainstate
 
 
 class TestOptaxOptimizer(unittest.TestCase):
     def test1(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.linear1 = bst.nn.Linear(2, 3)
-                self.linear2 = bst.nn.Linear(3, 4)
+                self.linear1 = brainstate.nn.Linear(2, 3)
+                self.linear2 = brainstate.nn.Linear(3, 4)
 
             def __call__(self, x):
                 return self.linear2(self.linear1(x))
 
-        x = bst.random.randn(1, 2)
+        x = brainstate.random.randn(1, 2)
         y = jax.numpy.ones((1, 4))
 
         model = Model()
         tx = optax.adam(1e-3)
-        optimizer = bst.optim.OptaxOptimizer(tx)
-        optimizer.register_trainable_weights(model.states(bst.ParamState))
+        optimizer = brainstate.optim.OptaxOptimizer(tx)
+        optimizer.register_trainable_weights(model.states(brainstate.ParamState))
 
         loss_fn = lambda: ((model(x) - y) ** 2).mean()
         prev_loss = loss_fn()
 
-        grads = bst.augment.grad(loss_fn, model.states(bst.ParamState))()
+        grads = brainstate.augment.grad(loss_fn, model.states(brainstate.ParamState))()
         optimizer.update(grads)
 
         new_loss = loss_fn()

brainstate/optim/_sgd_optimizer.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 
 # -*- coding: utf-8 -*-
-from __future__ import annotations
 
 import functools
 from typing import Union, Dict, Optional, Tuple, Any, TypeVar

brainstate/random/_rand_funs.py

@@ -15,7 +15,6 @@
 
 
 # -*- coding: utf-8 -*-
-from __future__ import annotations
 
 from typing import Optional