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

brainstate/functional/_activations_test.py

@@ -25,48 +25,48 @@ from absl.testing import parameterized
 from jax._src import test_util as jtu
 from jax.test_util import check_grads
 
-import brainstate as bst
+import brainstate
 
 
 class NNFunctionsTest(jtu.JaxTestCase):
     @jtu.skip_on_flag("jax_skip_slow_tests", True)
     def testSoftplusGrad(self):
-        check_grads(bst.functional.softplus, (1e-8,), order=4, )
+        check_grads(brainstate.functional.softplus, (1e-8,), order=4, )
 
     def testSoftplusGradZero(self):
-        check_grads(bst.functional.softplus, (0.,), order=1)
+        check_grads(brainstate.functional.softplus, (0.,), order=1)
 
     def testSoftplusGradInf(self):
-        self.assertAllClose(1., jax.grad(bst.functional.softplus)(float('inf')))
+        self.assertAllClose(1., jax.grad(brainstate.functional.softplus)(float('inf')))
 
     def testSoftplusGradNegInf(self):
-        check_grads(bst.functional.softplus, (-float('inf'),), order=1)
+        check_grads(brainstate.functional.softplus, (-float('inf'),), order=1)
 
     def testSoftplusGradNan(self):
-        check_grads(bst.functional.softplus, (float('nan'),), order=1)
+        check_grads(brainstate.functional.softplus, (float('nan'),), order=1)
 
     @parameterized.parameters([int, float] + jtu.dtypes.floating + jtu.dtypes.integer)
     def testSoftplusZero(self, dtype):
-        self.assertEqual(jnp.log(dtype(2)), bst.functional.softplus(dtype(0)))
+        self.assertEqual(jnp.log(dtype(2)), brainstate.functional.softplus(dtype(0)))
 
     def testSparseplusGradZero(self):
-        check_grads(bst.functional.sparse_plus, (-2.,), order=1)
+        check_grads(brainstate.functional.sparse_plus, (-2.,), order=1)
 
     def testSparseplusGrad(self):
-        check_grads(bst.functional.sparse_plus, (0.,), order=1)
+        check_grads(brainstate.functional.sparse_plus, (0.,), order=1)
 
     def testSparseplusAndSparseSigmoid(self):
         self.assertAllClose(
-            jax.grad(bst.functional.sparse_plus)(0.),
-            bst.functional.sparse_sigmoid(0.),
+            jax.grad(brainstate.functional.sparse_plus)(0.),
+            brainstate.functional.sparse_sigmoid(0.),
             check_dtypes=False)
         self.assertAllClose(
-            jax.grad(bst.functional.sparse_plus)(2.),
-            bst.functional.sparse_sigmoid(2.),
+            jax.grad(brainstate.functional.sparse_plus)(2.),
+            brainstate.functional.sparse_sigmoid(2.),
             check_dtypes=False)
         self.assertAllClose(
-            jax.grad(bst.functional.sparse_plus)(-2.),
-            bst.functional.sparse_sigmoid(-2.),
+            jax.grad(brainstate.functional.sparse_plus)(-2.),
+            brainstate.functional.sparse_sigmoid(-2.),
             check_dtypes=False)
 
     #   def testSquareplusGrad(self):
@@ -107,55 +107,55 @@ class NNFunctionsTest(jtu.JaxTestCase):
 
     @parameterized.parameters([float] + jtu.dtypes.floating)
     def testMishZero(self, dtype):
-        self.assertEqual(dtype(0), bst.functional.mish(dtype(0)))
+        self.assertEqual(dtype(0), brainstate.functional.mish(dtype(0)))
 
     def testReluGrad(self):
         rtol = None
-        check_grads(bst.functional.relu, (1.,), order=3, rtol=rtol)
-        check_grads(bst.functional.relu, (-1.,), order=3, rtol=rtol)
-        jaxpr = jax.make_jaxpr(jax.grad(bst.functional.relu))(0.)
+        check_grads(brainstate.functional.relu, (1.,), order=3, rtol=rtol)
+        check_grads(brainstate.functional.relu, (-1.,), order=3, rtol=rtol)
+        jaxpr = jax.make_jaxpr(jax.grad(brainstate.functional.relu))(0.)
         self.assertGreaterEqual(len(jaxpr.jaxpr.eqns), 2)
 
     def testRelu6Grad(self):
         rtol = None
-        check_grads(bst.functional.relu6, (1.,), order=3, rtol=rtol)
-        check_grads(bst.functional.relu6, (-1.,), order=3, rtol=rtol)
-        self.assertAllClose(jax.grad(bst.functional.relu6)(0.), 0., check_dtypes=False)
-        self.assertAllClose(jax.grad(bst.functional.relu6)(6.), 0., check_dtypes=False)
+        check_grads(brainstate.functional.relu6, (1.,), order=3, rtol=rtol)
+        check_grads(brainstate.functional.relu6, (-1.,), order=3, rtol=rtol)
+        self.assertAllClose(jax.grad(brainstate.functional.relu6)(0.), 0., check_dtypes=False)
+        self.assertAllClose(jax.grad(brainstate.functional.relu6)(6.), 0., check_dtypes=False)
 
     def testSoftplusValue(self):
-        val = bst.functional.softplus(89.)
+        val = brainstate.functional.softplus(89.)
         self.assertAllClose(val, 89., check_dtypes=False)
 
     def testSparseplusValue(self):
-        val = bst.functional.sparse_plus(89.)
+        val = brainstate.functional.sparse_plus(89.)
         self.assertAllClose(val, 89., check_dtypes=False)
 
     def testSparsesigmoidValue(self):
-        self.assertAllClose(bst.functional.sparse_sigmoid(-2.), 0., check_dtypes=False)
-        self.assertAllClose(bst.functional.sparse_sigmoid(2.), 1., check_dtypes=False)
-        self.assertAllClose(bst.functional.sparse_sigmoid(0.), .5, check_dtypes=False)
+        self.assertAllClose(brainstate.functional.sparse_sigmoid(-2.), 0., check_dtypes=False)
+        self.assertAllClose(brainstate.functional.sparse_sigmoid(2.), 1., check_dtypes=False)
+        self.assertAllClose(brainstate.functional.sparse_sigmoid(0.), .5, check_dtypes=False)
 
     #   def testSquareplusValue(self):
     #     val = bst.functional.squareplus(1e3)
     #     self.assertAllClose(val, 1e3, check_dtypes=False, atol=1e-3)
 
     def testMishValue(self):
-        val = bst.functional.mish(1e3)
+        val = brainstate.functional.mish(1e3)
         self.assertAllClose(val, 1e3, check_dtypes=False, atol=1e-3)
 
     def testEluValue(self):
-        val = bst.functional.elu(1e4)
+        val = brainstate.functional.elu(1e4)
         self.assertAllClose(val, 1e4, check_dtypes=False)
 
     def testGluValue(self):
-        val = bst.functional.glu(jnp.array([1.0, 0.0]), axis=0)
+        val = brainstate.functional.glu(jnp.array([1.0, 0.0]), axis=0)
         self.assertAllClose(val, jnp.array([0.5]))
 
     @parameterized.parameters(False, True)
     def testGeluIntType(self, approximate):
-        val_float = bst.functional.gelu(jnp.array(-1.0), approximate=approximate)
-        val_int = bst.functional.gelu(jnp.array(-1), approximate=approximate)
+        val_float = brainstate.functional.gelu(jnp.array(-1.0), approximate=approximate)
+        val_int = brainstate.functional.gelu(jnp.array(-1), approximate=approximate)
         self.assertAllClose(val_float, val_int)
 
     @parameterized.parameters(False, True)
@@ -166,19 +166,19 @@ class NNFunctionsTest(jtu.JaxTestCase):
         rng = jtu.rand_default(self.rng())
         args_maker = lambda: [rng((4, 5, 6), jnp.float32)]
         self._CheckAgainstNumpy(
-            gelu_reference, partial(bst.functional.gelu, approximate=approximate), args_maker,
+            gelu_reference, partial(brainstate.functional.gelu, approximate=approximate), args_maker,
             check_dtypes=False, tol=1e-3 if approximate else None)
 
     @parameterized.parameters(*itertools.product(
         (jnp.float32, jnp.bfloat16, jnp.float16),
-        (partial(bst.functional.gelu, approximate=False),
-         partial(bst.functional.gelu, approximate=True),
-         bst.functional.relu,
-         bst.functional.softplus,
-         bst.functional.sparse_plus,
-         bst.functional.sigmoid,
+        (partial(brainstate.functional.gelu, approximate=False),
+         partial(brainstate.functional.gelu, approximate=True),
+         brainstate.functional.relu,
+         brainstate.functional.softplus,
+         brainstate.functional.sparse_plus,
+         brainstate.functional.sigmoid,
          #  bst.functional.squareplus,
-         bst.functional.mish)))
+         brainstate.functional.mish)))
     def testDtypeMatchesInput(self, dtype, fn):
         x = jnp.zeros((), dtype=dtype)
         out = fn(x)
@@ -187,26 +187,26 @@ class NNFunctionsTest(jtu.JaxTestCase):
     def testEluMemory(self):
         # see https://github.com/google/jax/pull/1640
         with jax.enable_checks(False):  # With checks we materialize the array
-            jax.make_jaxpr(lambda: bst.functional.elu(jnp.ones((10 ** 12,))))  # don't oom
+            jax.make_jaxpr(lambda: brainstate.functional.elu(jnp.ones((10 ** 12,))))  # don't oom
 
     def testHardTanhMemory(self):
         # see https://github.com/google/jax/pull/1640
         with jax.enable_checks(False):  # With checks we materialize the array
-            jax.make_jaxpr(lambda: bst.functional.hard_tanh(jnp.ones((10 ** 12,))))  # don't oom
+            jax.make_jaxpr(lambda: brainstate.functional.hard_tanh(jnp.ones((10 ** 12,))))  # don't oom
 
-    @parameterized.parameters([bst.functional.softmax, bst.functional.log_softmax])
+    @parameterized.parameters([brainstate.functional.softmax, brainstate.functional.log_softmax])
     def testSoftmaxEmptyArray(self, fn):
         x = jnp.array([], dtype=float)
         self.assertArraysEqual(fn(x), x)
 
-    @parameterized.parameters([bst.functional.softmax, bst.functional.log_softmax])
+    @parameterized.parameters([brainstate.functional.softmax, brainstate.functional.log_softmax])
     def testSoftmaxEmptyMask(self, fn):
         x = jnp.array([5.5, 1.3, -4.2, 0.9])
         m = jnp.zeros_like(x, dtype=bool)
-        expected = jnp.full_like(x, 0.0 if fn is bst.functional.softmax else -jnp.inf)
+        expected = jnp.full_like(x, 0.0 if fn is brainstate.functional.softmax else -jnp.inf)
         self.assertArraysEqual(fn(x, where=m), expected)
 
-    @parameterized.parameters([bst.functional.softmax, bst.functional.log_softmax])
+    @parameterized.parameters([brainstate.functional.softmax, brainstate.functional.log_softmax])
     def testSoftmaxWhereMask(self, fn):
         x = jnp.array([5.5, 1.3, -4.2, 0.9])
         m = jnp.array([True, False, True, True])
@@ -214,10 +214,10 @@ class NNFunctionsTest(jtu.JaxTestCase):
         out = fn(x, where=m)
         self.assertAllClose(out[m], fn(x[m]))
 
-        probs = out if fn is bst.functional.softmax else jnp.exp(out)
+        probs = out if fn is brainstate.functional.softmax else jnp.exp(out)
         self.assertAllClose(probs.sum(), 1.0)
 
-    @parameterized.parameters([bst.functional.softmax, bst.functional.log_softmax])
+    @parameterized.parameters([brainstate.functional.softmax, brainstate.functional.log_softmax])
     def testSoftmaxWhereGrad(self, fn):
         # regression test for https://github.com/google/jax/issues/19490
         x = jnp.array([36., 10000.])
@@ -229,46 +229,46 @@ class NNFunctionsTest(jtu.JaxTestCase):
 
     def testSoftmaxGrad(self):
         x = jnp.array([5.5, 1.3, -4.2, 0.9])
-        jtu.check_grads(bst.functional.softmax, (x,), order=2, atol=5e-3)
+        jtu.check_grads(brainstate.functional.softmax, (x,), order=2, atol=5e-3)
 
     def testStandardizeWhereMask(self):
         x = jnp.array([5.5, 1.3, -4.2, 0.9])
         m = jnp.array([True, False, True, True])
         x_filtered = jnp.take(x, jnp.array([0, 2, 3]))
 
-        out_masked = jnp.take(bst.functional.standardize(x, where=m), jnp.array([0, 2, 3]))
-        out_filtered = bst.functional.standardize(x_filtered)
+        out_masked = jnp.take(brainstate.functional.standardize(x, where=m), jnp.array([0, 2, 3]))
+        out_filtered = brainstate.functional.standardize(x_filtered)
 
         self.assertAllClose(out_masked, out_filtered)
 
     def testOneHot(self):
-        actual = bst.functional.one_hot(jnp.array([0, 1, 2]), 3)
+        actual = brainstate.functional.one_hot(jnp.array([0, 1, 2]), 3)
         expected = jnp.array([[1., 0., 0.],
                               [0., 1., 0.],
                               [0., 0., 1.]])
         self.assertAllClose(actual, expected, check_dtypes=False)
 
-        actual = bst.functional.one_hot(jnp.array([1, 2, 0]), 3)
+        actual = brainstate.functional.one_hot(jnp.array([1, 2, 0]), 3)
         expected = jnp.array([[0., 1., 0.],
                               [0., 0., 1.],
                               [1., 0., 0.]])
         self.assertAllClose(actual, expected, check_dtypes=False)
 
     def testOneHotOutOfBound(self):
-        actual = bst.functional.one_hot(jnp.array([-1, 3]), 3)
+        actual = brainstate.functional.one_hot(jnp.array([-1, 3]), 3)
         expected = jnp.array([[0., 0., 0.],
                               [0., 0., 0.]])
         self.assertAllClose(actual, expected, check_dtypes=False)
 
     def testOneHotNonArrayInput(self):
-        actual = bst.functional.one_hot([0, 1, 2], 3)
+        actual = brainstate.functional.one_hot([0, 1, 2], 3)
         expected = jnp.array([[1., 0., 0.],
                               [0., 1., 0.],
                               [0., 0., 1.]])
         self.assertAllClose(actual, expected, check_dtypes=False)
 
     def testOneHotCustomDtype(self):
-        actual = bst.functional.one_hot(jnp.array([0, 1, 2]), 3, dtype=jnp.bool_)
+        actual = brainstate.functional.one_hot(jnp.array([0, 1, 2]), 3, dtype=jnp.bool_)
         expected = jnp.array([[True, False, False],
                               [False, True, False],
                               [False, False, True]])
@@ -279,14 +279,14 @@ class NNFunctionsTest(jtu.JaxTestCase):
                               [0., 0., 1.],
                               [1., 0., 0.]]).T
 
-        actual = bst.functional.one_hot(jnp.array([1, 2, 0]), 3, axis=0)
+        actual = brainstate.functional.one_hot(jnp.array([1, 2, 0]), 3, axis=0)
         self.assertAllClose(actual, expected, check_dtypes=False)
 
-        actual = bst.functional.one_hot(jnp.array([1, 2, 0]), 3, axis=-2)
+        actual = brainstate.functional.one_hot(jnp.array([1, 2, 0]), 3, axis=-2)
         self.assertAllClose(actual, expected, check_dtypes=False)
 
     def testTanhExists(self):
-        print(bst.functional.tanh)  # doesn't crash
+        print(brainstate.functional.tanh)  # doesn't crash
 
     def testCustomJVPLeak(self):
         # https://github.com/google/jax/issues/8171
@@ -295,7 +295,7 @@ class NNFunctionsTest(jtu.JaxTestCase):
             a = jnp.array(1.)
 
             def f(hx, _):
-                hx = bst.functional.sigmoid(hx + a)
+                hx = brainstate.functional.sigmoid(hx + a)
                 return hx, None
 
             hx = jnp.array(0.)

brainstate/functional/_normalization.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 from typing import Optional, Union
 
 import brainunit as u

brainstate/functional/_others.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 from functools import partial
 
 import jax

brainstate/functional/_spikes.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 __all__ = [
     'spike_bitwise_or',

brainstate/graph/_graph_node.py

@@ -15,8 +15,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from abc import ABCMeta
 from copy import deepcopy
 from typing import Any, Callable, Type, TypeVar, Tuple, TYPE_CHECKING, Mapping, Iterator, Sequence
@@ -210,7 +208,7 @@ class List(Node):
     def __len__(self):
         return len(vars(self))
 
-    def __add__(self, other: Sequence[A]) -> List[A]:
+    def __add__(self, other: Sequence[A]) -> 'List[A]':
         return List(list(self) + list(other))
 
     def append(self, value):

brainstate/graph/_graph_node_test.py

@@ -13,63 +13,61 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import unittest
 
-import brainstate as bst
+import brainstate
 
 
 class TestSequential(unittest.TestCase):
     def test1(self):
-        s = bst.graph.Sequential(bst.nn.Linear(1, 2),
-                                 bst.nn.Linear(2, 3))
-        graphdef, states = bst.graph.treefy_split(s)
+        s = brainstate.graph.Sequential(brainstate.nn.Linear(1, 2),
+                                        brainstate.nn.Linear(2, 3))
+        graphdef, states = brainstate.graph.treefy_split(s)
         print(states)
         self.assertTrue(len(states.to_flat()) == 2)
 
 
 class TestStateRetrieve(unittest.TestCase):
     def test_list_of_states_1(self):
-        class Model(bst.graph.Node):
+        class Model(brainstate.graph.Node):
             def __init__(self):
                 self.a = [1, 2, 3]
-                self.b = [bst.State(1), bst.State(2), bst.State(3)]
+                self.b = [brainstate.State(1), brainstate.State(2), brainstate.State(3)]
 
         m = Model()
-        graphdef, states = bst.graph.treefy_split(m)
+        graphdef, states = brainstate.graph.treefy_split(m)
         print(states.to_flat())
         self.assertTrue(len(states.to_flat()) == 3)
 
     def test_list_of_states_2(self):
-        class Model(bst.graph.Node):
+        class Model(brainstate.graph.Node):
             def __init__(self):
                 self.a = [1, 2, 3]
-                self.b = [bst.State(1), [bst.State(2), bst.State(3)]]
+                self.b = [brainstate.State(1), [brainstate.State(2), brainstate.State(3)]]
 
         m = Model()
-        graphdef, states = bst.graph.treefy_split(m)
+        graphdef, states = brainstate.graph.treefy_split(m)
         print(states.to_flat())
         self.assertTrue(len(states.to_flat()) == 3)
 
     def test_list_of_node_1(self):
-        class Model(bst.graph.Node):
+        class Model(brainstate.graph.Node):
             def __init__(self):
                 self.a = [1, 2, 3]
-                self.b = [bst.nn.Linear(1, 2), bst.nn.Linear(2, 3)]
+                self.b = [brainstate.nn.Linear(1, 2), brainstate.nn.Linear(2, 3)]
 
         m = Model()
-        graphdef, states = bst.graph.treefy_split(m)
+        graphdef, states = brainstate.graph.treefy_split(m)
         print(states.to_flat())
         self.assertTrue(len(states.to_flat()) == 2)
 
     def test_list_of_node_2(self):
-        class Model(bst.graph.Node):
+        class Model(brainstate.graph.Node):
             def __init__(self):
                 self.a = [1, 2, 3]
-                self.b = [bst.nn.Linear(1, 2), [bst.nn.Linear(2, 3)], (bst.nn.Linear(3, 4), bst.nn.Linear(4, 5))]
+                self.b = [brainstate.nn.Linear(1, 2), [brainstate.nn.Linear(2, 3)], (brainstate.nn.Linear(3, 4), brainstate.nn.Linear(4, 5))]
 
         m = Model()
-        graphdef, states = bst.graph.treefy_split(m)
+        graphdef, states = brainstate.graph.treefy_split(m)
         print(states.to_flat())
         self.assertTrue(len(states.to_flat()) == 4)

brainstate/graph/_graph_operation.py

@@ -18,8 +18,10 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import (Any, Callable, Generic, Iterable, Iterator, Mapping, MutableMapping,
-                    Sequence, Type, TypeVar, Union, Hashable, Tuple, Dict, Optional, overload)
+from typing import (
+    Any, Callable, Generic, Iterable, Iterator, Mapping, MutableMapping,
+    Sequence, Type, TypeVar, Union, Hashable, Tuple, Dict, Optional, overload
+)
 
 import jax
 import numpy as np