brainstate 0.1.5__py2.py3-none-any.whl → 0.1.6__py2.py3-none-any.whl

brainstate/surrogate.py

@@ -106,11 +106,11 @@ class Surrogate(PrettyObject):
     Examples
     --------
 
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> import brainstate.nn as nn
     >>> import jax.numpy as jnp
 
-    >>> class MySurrogate(bst.surrogate.Surrogate):
+    >>> class MySurrogate(brainstate.surrogate.Surrogate):
     ...   def __init__(self, alpha=1.):
     ...     super().__init__()
     ...     self.alpha = alpha
@@ -236,11 +236,11 @@ def sigmoid(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-2, 2, 1000)
        >>> for alpha in [1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.sigmoid)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.sigmoid)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -355,11 +355,11 @@ def piecewise_quadratic(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.piecewise_quadratic)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.piecewise_quadratic)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -522,11 +522,11 @@ def piecewise_exp(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.piecewise_exp)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.piecewise_exp)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -621,11 +621,11 @@ def soft_sign(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.soft_sign)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.soft_sign)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -706,11 +706,11 @@ def arctan(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.arctan)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.arctan)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -804,11 +804,11 @@ def nonzero_sign_log(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.nonzero_sign_log)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.nonzero_sign_log)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -893,11 +893,11 @@ def erf(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.nonzero_sign_log)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.nonzero_sign_log)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -1000,12 +1000,12 @@ def piecewise_leaky_relu(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for c in [0.01, 0.05, 0.1]:
        >>>   for w in [1., 2.]:
-       >>>     grads1 = bst.augment.vector_grad(bst.surrogate.piecewise_leaky_relu)(xs, c=c, w=w)
+       >>>     grads1 = brainstate.augment.vector_grad(brainstate.surrogate.piecewise_leaky_relu)(xs, c=c, w=w)
        >>>     plt.plot(xs, grads1, label=f'x={c}, w={w}')
        >>> plt.legend()
        >>> plt.show()
@@ -1113,12 +1113,12 @@ def squarewave_fourier_series(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for n in [2, 4, 8]:
-       >>>   f = bst.surrogate.SquarewaveFourierSeries(n=n)
-       >>>   grads1 = bst.augment.vector_grad(f)(xs)
+       >>>   f = brainstate.surrogate.SquarewaveFourierSeries(n=n)
+       >>>   grads1 = brainstate.augment.vector_grad(f)(xs)
        >>>   plt.plot(xs, grads1, label=f'n={n}')
        >>> plt.legend()
        >>> plt.show()
@@ -1214,12 +1214,12 @@ def s2nn(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
-       >>> grads = bst.augment.vector_grad(bst.surrogate.s2nn)(xs, 4., 1.)
+       >>> grads = brainstate.augment.vector_grad(brainstate.surrogate.s2nn)(xs, 4., 1.)
        >>> plt.plot(xs, grads, label=r'$\alpha=4, \beta=1$')
-       >>> grads = bst.augment.vector_grad(bst.surrogate.s2nn)(xs, 8., 2.)
+       >>> grads = brainstate.augment.vector_grad(brainstate.surrogate.s2nn)(xs, 8., 2.)
        >>> plt.plot(xs, grads, label=r'$\alpha=8, \beta=2$')
        >>> plt.legend()
        >>> plt.show()
@@ -1315,11 +1315,11 @@ def q_pseudo_spike(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.q_pseudo_spike)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.q_pseudo_spike)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha=$' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -1413,10 +1413,10 @@ def leaky_relu(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
-       >>> grads = bst.augment.vector_grad(bst.surrogate.leaky_relu)(xs, 0., 1.)
+       >>> grads = brainstate.augment.vector_grad(brainstate.surrogate.leaky_relu)(xs, 0., 1.)
        >>> plt.plot(xs, grads, label=r'$\alpha=0., \beta=1.$')
        >>> plt.legend()
        >>> plt.show()
@@ -1517,10 +1517,10 @@ def log_tailed_relu(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
-       >>> grads = bst.augment.vector_grad(bst.surrogate.leaky_relu)(xs, 0., 1.)
+       >>> grads = brainstate.augment.vector_grad(brainstate.surrogate.leaky_relu)(xs, 0., 1.)
        >>> plt.plot(xs, grads, label=r'$\alpha=0., \beta=1.$')
        >>> plt.legend()
        >>> plt.show()
@@ -1596,12 +1596,12 @@ def relu_grad(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for s in [0.5, 1.]:
        >>>   for w in [1, 2.]:
-       >>>     grads = bst.augment.vector_grad(bst.surrogate.relu_grad)(xs, s, w)
+       >>>     grads = brainstate.augment.vector_grad(brainstate.surrogate.relu_grad)(xs, s, w)
        >>>     plt.plot(xs, grads, label=r'$\alpha=$' + f'{s}, width={w}')
        >>> plt.legend()
        >>> plt.show()
@@ -1678,11 +1678,11 @@ def gaussian_grad(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for s in [0.5, 1., 2.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.gaussian_grad)(xs, s, 0.5)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.gaussian_grad)(xs, s, 0.5)
        >>>   plt.plot(xs, grads, label=r'$\alpha=0.5, \sigma=$' + str(s))
        >>> plt.legend()
        >>> plt.show()
@@ -1773,10 +1773,10 @@ def multi_gaussian_grad(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
-       >>> grads = bst.augment.vector_grad(bst.surrogate.multi_gaussian_grad)(xs)
+       >>> grads = brainstate.augment.vector_grad(brainstate.surrogate.multi_gaussian_grad)(xs)
        >>> plt.plot(xs, grads)
        >>> plt.show()
 
@@ -1855,11 +1855,11 @@ def inv_square_grad(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-1, 1, 1000)
        >>> for alpha in [1., 10., 100.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.inv_square_grad)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.inv_square_grad)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -1929,11 +1929,11 @@ def slayer_grad(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.slayer_grad)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.slayer_grad)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()

brainstate/util/pretty_pytree.py

@@ -373,19 +373,19 @@ class NestedDict(PrettyDict):
 
         Example usage::
 
-          >>> import brainstate as bst
+          >>> import brainstate as brainstate
 
-          >>> class Model(bst.nn.Module):
+          >>> class Model(brainstate.nn.Module):
           ...   def __init__(self):
           ...     super().__init__()
-          ...     self.batchnorm = bst.nn.BatchNorm1d([10, 3])
-          ...     self.linear = bst.nn.Linear(2, 3)
+          ...     self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
+          ...     self.linear = brainstate.nn.Linear(2, 3)
           ...   def __call__(self, x):
           ...     return self.linear(self.batchnorm(x))
 
           >>> model = Model()
-          >>> state_map = bst.graph.treefy_states(model)
-          >>> param, others = state_map.treefy_split(bst.ParamState, ...)
+          >>> state_map = brainstate.graph.treefy_states(model)
+          >>> param, others = state_map.treefy_split(brainstate.ParamState, ...)
 
         Arguments:
           first: The first filter
@@ -495,14 +495,14 @@ class FlattedDict(PrettyDict):
 
     Example usage::
 
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax.numpy as jnp
       >>>
-      >>> class Model(bst.nn.Module):
+      >>> class Model(brainstate.nn.Module):
       ...   def __init__(self):
       ...     super().__init__()
-      ...     self.batchnorm = bst.nn.BatchNorm1d([10, 3])
-      ...     self.linear = bst.nn.Linear(2, 3)
+      ...     self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
+      ...     self.linear = brainstate.nn.Linear(2, 3)
       ...   def __call__(self, x):
       ...     return self.linear(self.batchnorm(x))
       >>>

brainstate/util/{_pretty_pytree_test.py → pretty_pytree_test.py}

@@ -13,31 +13,30 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
 import jax
 from absl.testing import absltest
 
-import brainstate as bst
+import brainstate
 
 
 class TestNestedMapping(absltest.TestCase):
     def test_create_state(self):
-        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+        state = brainstate.util.NestedDict({'a': brainstate.ParamState(1), 'b': {'c': brainstate.ParamState(2)}})
 
         assert state['a'].value == 1
         assert state['b']['c'].value == 2
 
     def test_get_attr(self):
-        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+        state = brainstate.util.NestedDict({'a': brainstate.ParamState(1), 'b': {'c': brainstate.ParamState(2)}})
 
         assert state.a.value == 1
         assert state.b['c'].value == 2
 
     def test_set_attr(self):
-        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+        state = brainstate.util.NestedDict({'a': brainstate.ParamState(1), 'b': {'c': brainstate.ParamState(2)}})
 
         state.a.value = 3
         state.b['c'].value = 4
@@ -46,36 +45,36 @@ class TestNestedMapping(absltest.TestCase):
         assert state['b']['c'].value == 4
 
     def test_set_attr_variables(self):
-        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+        state = brainstate.util.NestedDict({'a': brainstate.ParamState(1), 'b': {'c': brainstate.ParamState(2)}})
 
         state.a.value = 3
         state.b['c'].value = 4
 
-        assert isinstance(state.a, bst.ParamState)
+        assert isinstance(state.a, brainstate.ParamState)
         assert state.a.value == 3
-        assert isinstance(state.b['c'], bst.ParamState)
+        assert isinstance(state.b['c'], brainstate.ParamState)
         assert state.b['c'].value == 4
 
     def test_add_nested_attr(self):
-        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
-        state.b['d'] = bst.ParamState(5)
+        state = brainstate.util.NestedDict({'a': brainstate.ParamState(1), 'b': {'c': brainstate.ParamState(2)}})
+        state.b['d'] = brainstate.ParamState(5)
 
         assert state['b']['d'].value == 5
 
     def test_delete_nested_attr(self):
-        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+        state = brainstate.util.NestedDict({'a': brainstate.ParamState(1), 'b': {'c': brainstate.ParamState(2)}})
         del state['b']['c']
 
         assert 'c' not in state['b']
 
     def test_integer_access(self):
-        class Foo(bst.nn.Module):
+        class Foo(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.layers = [bst.nn.Linear(1, 2), bst.nn.Linear(2, 3)]
+                self.layers = [brainstate.nn.Linear(1, 2), brainstate.nn.Linear(2, 3)]
 
         module = Foo()
-        state_refs = bst.graph.treefy_states(module)
+        state_refs = brainstate.graph.treefy_states(module)
 
         assert module.layers[0].weight.value['weight'].shape == (1, 2)
         assert state_refs.layers[0]['weight'].value['weight'].shape == (1, 2)
@@ -83,8 +82,8 @@ class TestNestedMapping(absltest.TestCase):
         assert state_refs.layers[1]['weight'].value['weight'].shape == (2, 3)
 
     def test_pure_dict(self):
-        module = bst.nn.Linear(4, 5)
-        state_map = bst.graph.treefy_states(module)
+        module = brainstate.nn.Linear(4, 5)
+        state_map = brainstate.graph.treefy_states(module)
         pure_dict = state_map.to_pure_dict()
         assert isinstance(pure_dict, dict)
         assert isinstance(pure_dict['weight'].value['weight'], jax.Array)
@@ -93,27 +92,27 @@ class TestNestedMapping(absltest.TestCase):
 
 class TestSplit(unittest.TestCase):
     def test_split(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.batchnorm = bst.nn.BatchNorm1d([10, 3])
-                self.linear = bst.nn.Linear([10, 3], [10, 4])
+                self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
+                self.linear = brainstate.nn.Linear([10, 3], [10, 4])
 
             def __call__(self, x):
                 return self.linear(self.batchnorm(x))
 
-        with bst.environ.context(fit=True):
+        with brainstate.environ.context(fit=True):
             model = Model()
-            x = bst.random.randn(1, 10, 3)
+            x = brainstate.random.randn(1, 10, 3)
             y = model(x)
             self.assertEqual(y.shape, (1, 10, 4))
 
-        state_map = bst.graph.treefy_states(model)
+        state_map = brainstate.graph.treefy_states(model)
 
         with self.assertRaises(ValueError):
-            params, others = state_map.split(bst.ParamState)
+            params, others = state_map.split(brainstate.ParamState)
 
-        params, others = state_map.split(bst.ParamState, ...)
+        params, others = state_map.split(brainstate.ParamState, ...)
         print()
         print(params)
         print(others)
@@ -124,37 +123,37 @@ class TestSplit(unittest.TestCase):
 
 class TestStateMap2(unittest.TestCase):
     def test1(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.batchnorm = bst.nn.BatchNorm1d([10, 3])
-                self.linear = bst.nn.Linear([10, 3], [10, 4])
+                self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
+                self.linear = brainstate.nn.Linear([10, 3], [10, 4])
 
             def __call__(self, x):
                 return self.linear(self.batchnorm(x))
 
-        with bst.environ.context(fit=True):
+        with brainstate.environ.context(fit=True):
             model = Model()
-            state_map = bst.graph.treefy_states(model).to_flat()
-            state_map = bst.util.NestedDict(state_map)
+            state_map = brainstate.graph.treefy_states(model).to_flat()
+            state_map = brainstate.util.NestedDict(state_map)
 
 
 class TestFlattedMapping(unittest.TestCase):
     def test1(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.batchnorm = bst.nn.BatchNorm1d([10, 3])
-                self.linear = bst.nn.Linear([10, 3], [10, 4])
+                self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
+                self.linear = brainstate.nn.Linear([10, 3], [10, 4])
 
             def __call__(self, x):
                 return self.linear(self.batchnorm(x))
 
         model = Model()
         # print(model.states())
-        # print(bst.graph.states(model))
-        self.assertTrue(model.states() == bst.graph.states(model))
+        # print(brainstate.graph.states(model))
+        self.assertTrue(model.states() == brainstate.graph.states(model))
 
         print(model.nodes())
-        # print(bst.graph.nodes(model))
-        self.assertTrue(model.nodes() == bst.graph.nodes(model))
+        # print(brainstate.graph.nodes(model))
+        self.assertTrue(model.nodes() == brainstate.graph.nodes(model))

brainstate/util/struct.py

@@ -56,16 +56,16 @@ def dataclass(clz: T, **kwargs) -> T:
     The ``dataclass`` decorator makes it easy to define custom classes that can be
     passed safely to Jax. For example::
 
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax
       >>> from typing import Any, Callable
 
-      >>> @bst.util.dataclass
+      >>> @brainstate.util.dataclass
       ... class Model:
       ...   params: Any
       ...   # use pytree_node=False to indicate an attribute should not be touched
       ...   # by Jax transformations.
-      ...   apply_fn: Callable = bst.util.field(pytree_node=False)
+      ...   apply_fn: Callable = brainstate.util.field(pytree_node=False)
 
       ...   def __apply__(self, *args):
       ...     return self.apply_fn(*args)
@@ -97,7 +97,7 @@ def dataclass(clz: T, **kwargs) -> T:
     This way the simple constructor used by ``jax.tree_util`` is
     preserved. Consider the following example::
 
-      >>> @bst.util.dataclass
+      >>> @brainstate.util.dataclass
       ... class DirectionAndScaleKernel:
       ...   direction: jax.Array
       ...   scale: jax.Array
@@ -189,15 +189,15 @@ class PyTreeNode:
 
     Example::
 
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax
       >>> from typing import Any, Callable
 
-      >>> class Model(bst.util.PyTreeNode):
+      >>> class Model(brainstate.util.PyTreeNode):
       ...   params: Any
       ...   # use pytree_node=False to indicate an attribute should not be touched
       ...   # by Jax transformations.
-      ...   apply_fn: Callable = bst.util.field(pytree_node=False)
+      ...   apply_fn: Callable = brainstate.util.field(pytree_node=False)
 
       ...   def __apply__(self, *args):
       ...     return self.apply_fn(*args)

{brainstate-0.1.5.dist-info → brainstate-0.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.5
+Version: 0.1.6
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers
@@ -31,22 +31,22 @@ License-File: LICENSE
 Requires-Dist: jax
 Requires-Dist: jaxlib
 Requires-Dist: numpy
-Requires-Dist: brainunit (>=0.0.4)
+Requires-Dist: brainunit>=0.0.4
 Requires-Dist: brainevent
 Provides-Extra: cpu
-Requires-Dist: jax[cpu] ; extra == 'cpu'
-Requires-Dist: brainunit[cpu] ; extra == 'cpu'
-Requires-Dist: brainevent[cpu] ; extra == 'cpu'
+Requires-Dist: jax[cpu]; extra == "cpu"
+Requires-Dist: brainunit[cpu]; extra == "cpu"
+Requires-Dist: brainevent[cpu]; extra == "cpu"
 Provides-Extra: cuda12
-Requires-Dist: jax[cuda12] ; extra == 'cuda12'
-Requires-Dist: brainunit[cuda12] ; extra == 'cuda12'
-Requires-Dist: brainevent[cuda12] ; extra == 'cuda12'
+Requires-Dist: jax[cuda12]; extra == "cuda12"
+Requires-Dist: brainunit[cuda12]; extra == "cuda12"
+Requires-Dist: brainevent[cuda12]; extra == "cuda12"
 Provides-Extra: testing
-Requires-Dist: pytest ; extra == 'testing'
+Requires-Dist: pytest; extra == "testing"
 Provides-Extra: tpu
-Requires-Dist: jax[tpu] ; extra == 'tpu'
-Requires-Dist: brainunit[tpu] ; extra == 'tpu'
-Requires-Dist: brainevent[tpu] ; extra == 'tpu'
+Requires-Dist: jax[tpu]; extra == "tpu"
+Requires-Dist: brainunit[tpu]; extra == "tpu"
+Requires-Dist: brainevent[tpu]; extra == "tpu"
 
 
 # A ``State``-based Transformation System for Program Compilation and Augmentation