Fix spike return values and update file paths in simulation scripts

Author: chaoming0625

brainpy/_src/math/object_transform/variables.py

@@ -1,16 +1,16 @@
 from typing import Optional, Any, List, Callable, Sequence, Union, Dict, Tuple
 
-import brainstate
 import jax
 import numpy as np
-from brainstate._state import record_state_value_read, record_state_value_write
 from jax import numpy as jnp
 from jax.dtypes import canonicalize_dtype
 from jax.tree_util import register_pytree_node_class
 
+import brainstate
 from brainpy._src.math.ndarray import BaseArray
 from brainpy._src.math.sharding import BATCH_AXIS
 from brainpy.errors import MathError
+from brainstate._state import record_state_value_read, record_state_value_write
 
 __all__ = [
     'Variable',
@@ -313,8 +313,12 @@ def batch_size(self) -> Optional[int]:
     def batch_size(self, val):
         raise ValueError(f'Cannot set "batch_size" manually.')
 
+    def _ensure_value_exists(self):
+        pass
+
     @property
     def value(self):
+        self._ensure_value_exists()
         record_state_value_read(self)
         return self._read_value()
 

brainpy/_src/math/random.py

@@ -555,16 +555,6 @@ def _ensure_value_exists(self):
                         seed_or_key = np.random.randint(0, 100000, 2, dtype=np.uint32)
                         self._value = seed_or_key
 
-    @property
-    def value(self):
-        self._ensure_value_exists()
-        record_state_value_read(self)
-        return self._read_value()
-
-    # def check_if_deleted(self):
-    #   if self.value.is_deleted():
-    #     self.seed()
-
     def split_key(self):
         """Create a new seed from the current seed.
         """

brainpy/_src/transform.py

@@ -3,8 +3,8 @@
 import functools
 from typing import Union, Optional, Dict, Sequence
 
+import jax
 import jax.numpy as jnp
-from jax.tree_util import tree_flatten, tree_unflatten, tree_map
 
 from brainpy import tools, math as bm
 from brainpy._src.context import share
@@ -159,7 +159,7 @@ def __init__(
         self.no_state = no_state
         self.out_vars = out_vars
         if out_vars is not None:
-            out_vars, _ = tree_flatten(out_vars, is_leaf=lambda s: isinstance(s, bm.Variable))
+            out_vars, _ = jax.tree.flatten(out_vars, is_leaf=lambda s: isinstance(s, bm.Variable))
             for v in out_vars:
                 if not isinstance(v, bm.Variable):
                     raise TypeError('out_vars must be a PyTree of Variable.')
@@ -198,7 +198,7 @@ def __call__(
                            'Input should be a Array PyTree with the shape '
                            'of (B, T, ...) or (T, B, ...) with `data_first_axis="T"`, '
                            'where B the batch size and T the time length.')
-            xs, tree = tree_flatten(duration_or_xs, lambda a: isinstance(a, bm.BaseArray))
+            xs, tree = jax.tree.flatten(duration_or_xs, lambda a: isinstance(a, bm.BaseArray))
             if self.target.mode.is_child_of(bm.BatchingMode):
                 b_idx, t_idx = (1, 0) if self.data_first_axis == 'T' else (0, 1)
 
@@ -209,26 +209,26 @@ def __call__(
                 if len(batch) != 1:
                     raise ValueError('\n'
                                      'Input should be a Array PyTree with the same batch dimension. '
-                                     f'but we got {tree_unflatten(tree, batch)}.')
+                                     f'but we got {jax.tree.unflatten(tree, batch)}.')
                 try:
                     length = tuple(set([x.shape[t_idx] for x in xs]))
                 except (AttributeError, IndexError) as e:
                     raise ValueError(inp_err_msg) from e
                 if len(batch) != 1:
                     raise ValueError('\n'
                                      'Input should be a Array PyTree with the same batch size. '
-                                     f'but we got {tree_unflatten(tree, batch)}.')
+                                     f'but we got {jax.tree.unflatten(tree, batch)}.')
                 if len(length) != 1:
                     raise ValueError('\n'
                                      'Input should be a Array PyTree with the same time length. '
-                                     f'but we got {tree_unflatten(tree, length)}.')
+                                     f'but we got {jax.tree.unflatten(tree, length)}.')
 
                 if self.no_state:
                     xs = [bm.reshape(x, (length[0] * batch[0],) + x.shape[2:]) for x in xs]
                 else:
                     if self.data_first_axis == 'B':
                         xs = [jnp.moveaxis(x, 0, 1) for x in xs]
-                xs = tree_unflatten(tree, xs)
+                xs = jax.tree.unflatten(tree, xs)
                 origin_shape = (length[0], batch[0]) if self.data_first_axis == 'T' else (batch[0], length[0])
 
             else:
@@ -240,15 +240,15 @@ def __call__(
                 if len(length) != 1:
                     raise ValueError('\n'
                                      'Input should be a Array PyTree with the same time length. '
-                                     f'but we got {tree_unflatten(tree, length)}.')
-                xs = tree_unflatten(tree, xs)
+                                     f'but we got {jax.tree.unflatten(tree, length)}.')
+                xs = jax.tree.unflatten(tree, xs)
                 origin_shape = (length[0],)
 
             # computation
             if self.no_state:
                 share.save(**self.shared_arg)
                 outputs = self._run(self.shared_arg, dict(), xs)
-                results = tree_map(lambda a: jnp.reshape(a, origin_shape + a.shape[1:]), outputs)
+                results = jax.tree.map(lambda a: jnp.reshape(a, origin_shape + a.shape[1:]), outputs)
                 if self.i0 is not None:
                     self.i0 += length[0]
                 if self.t0 is not None:
@@ -263,6 +263,7 @@ def __call__(
                     shared['i'] = jnp.arange(0, length[0]) + self.i0.value
 
         assert not self.no_state
+        xs = jax.tree.map(lambda x: x.value if isinstance(x, bm.Variable) else x, xs, is_leaf=lambda x: isinstance(x, bm.Variable))
         results = bm.for_loop(functools.partial(self._run, self.shared_arg),
                               (shared, xs),
                               jit=self.jit,
@@ -283,6 +284,6 @@ def _run(self, static_sh, dyn_sh, x):
         share.save(**static_sh, **dyn_sh)
         outs = self.target(x)
         if self.out_vars is not None:
-            outs = (outs, tree_map(bm.as_jax, self.out_vars))
+            outs = (outs, jax.tree.map(bm.as_jax, self.out_vars, is_leaf=lambda x: isinstance(x, bm.Variable)))
         clear_input(self.target)
         return outs

examples/dynamics_simulation/COBA_parallel.py

@@ -72,7 +72,7 @@ def update(self, inp):
             self.I2I()
             self.E(inp)
             self.I(inp)
-            return self.E.spike
+            return self.E.spike.value
 
     model = EINet()
     indices = bm.arange(1000)
@@ -150,7 +150,7 @@ def update(self, inp):
             self.I2I()
             self.E(inp)
             self.I(inp)
-            return self.E.spike
+            return self.E.spike.value
 
     model = EINet()
     indices = bm.arange(1000)
@@ -199,7 +199,7 @@ def update(self, inp):
             self.I2I()
             self.E(inp)
             self.I(inp)
-            return self.E.spike
+            return self.E.spike.value
 
     model = EINet()
     indices = bm.arange(1000)
@@ -241,9 +241,9 @@ def update(self, input):
 
 
 if __name__ == '__main__':
-    model1()
-    # model2()
-    # model3()
-    # model4()
-    # model5()
-    # vanalla_proj()
+    # model1()
+    model2()
+    model3()
+    model4()
+    model5()
+    vanalla_proj()