Fix STDP weight update logic and improve dataset loading paths

Author: chaoming0625

brainpy/_src/dnn/linear.py

@@ -217,7 +217,7 @@ def stdp_update(
         if on_post is not None:
             spike = on_post['spike']
             trace = on_post['trace']
-            self.W.value = dense_on_post(self.W.value, spike, trace, w_min, w_max)
+            self.W.value = dense_on_post(self.W.value, trace, spike, w_min, w_max)
 
 
 Linear = Dense
@@ -303,18 +303,16 @@ def stdp_update(
         w_min: numbers.Number = None,
         w_max: numbers.Number = None
     ):
-        if isinstance(self.weight, float):
-            raise ValueError(f'Cannot update the weight of a constant node.')
         if not isinstance(self.weight, bm.Variable):
-            self.tracing_variable('weight', self.weight, self.weight.shape)
+            raise ValueError(f'When using STDP to update synaptic weights, the weight must be a variable.')
         if on_pre is not None:
             spike = on_pre['spike']
             trace = on_pre['trace']
             self.weight.value = dense_on_pre(self.weight.value, spike, trace, w_min, w_max)
         if on_post is not None:
             spike = on_post['spike']
             trace = on_post['trace']
-            self.weight.value = dense_on_post(self.weight.value, spike, trace, w_min, w_max)
+            self.weight.value = dense_on_post(self.weight.value, trace, spike, w_min, w_max)
 
 
 class OneToOne(Layer, SupportSTDP):
@@ -442,7 +440,7 @@ def stdp_update(
         if on_post is not None:
             spike = on_post['spike']
             trace = on_post['trace']
-            self.weight.value = dense_on_post(self.weight.value, spike, trace, w_min, w_max)
+            self.weight.value = dense_on_post(self.weight.value, trace, spike, w_min, w_max)
 
 
 class _CSRLayer(Layer, SupportSTDP):
@@ -503,7 +501,7 @@ def stdp_update(
             trace = on_post['trace']
             self.weight.value = csr2csc_on_post(
                 self.weight.value, self._pre_ids, self._post_indptr,
-                self.w_indices, spike, trace, w_min, w_max,
+                self.w_indices, trace, spike, w_min, w_max,
                 shape=(trace.shape[0], spike.shape[0]),
             )
 

brainpy/_src/math/object_transform/variables.py

@@ -349,63 +349,63 @@ def value(self, v):
         self._been_writen = True  # set the flag
         self._write_value(v)  # write the value
 
-    def tree_flatten(self):
-        """Flattens this variable.
-
-        Returns:
-          A pair where the first element is a list of leaf values
-          and the second element is a treedef representing the
-          structure of the flattened tree.
-        """
-        return (self._value,), None
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, flat_contents):
-        """Reconstructs a variable from the aux_data and the leaves.
-
-        Args:
-          aux_data:
-          flat_contents:
-
-        Returns:
-          The variable.
-        """
-        return cls(*flat_contents)
-
-    def clone(self) -> 'Variable':
-        """Clone the variable. """
-        r = type(self)(jnp.array(self.value, copy=True), batch_axis=self.batch_axis)
-        return r
-
-    def __eq__(self, other):
-        """Override State's __eq__ to use BaseArray behavior for element-wise comparison."""
-        from brainpy._src.math.ndarray import _check_input_array, _return
-        return _return(self.value == _check_input_array(other))
-
-    def __ne__(self, other):
-        """Override State's __ne__ to use BaseArray behavior for element-wise comparison."""
-        from brainpy._src.math.ndarray import _check_input_array, _return
-        return _return(self.value != _check_input_array(other))
-
-    def __lt__(self, other):
-        """Override State's __lt__ to use BaseArray behavior for element-wise comparison."""
-        from brainpy._src.math.ndarray import _check_input_array, _return
-        return _return(self.value < _check_input_array(other))
-
-    def __le__(self, other):
-        """Override State's __le__ to use BaseArray behavior for element-wise comparison."""
-        from brainpy._src.math.ndarray import _check_input_array, _return
-        return _return(self.value <= _check_input_array(other))
-
-    def __gt__(self, other):
-        """Override State's __gt__ to use BaseArray behavior for element-wise comparison."""
-        from brainpy._src.math.ndarray import _check_input_array, _return
-        return _return(self.value > _check_input_array(other))
-
-    def __ge__(self, other):
-        """Override State's __ge__ to use BaseArray behavior for element-wise comparison."""
-        from brainpy._src.math.ndarray import _check_input_array, _return
-        return _return(self.value >= _check_input_array(other))
+    # def tree_flatten(self):
+    #     """Flattens this variable.
+    #
+    #     Returns:
+    #       A pair where the first element is a list of leaf values
+    #       and the second element is a treedef representing the
+    #       structure of the flattened tree.
+    #     """
+    #     return (self._value,), None
+    #
+    # @classmethod
+    # def tree_unflatten(cls, aux_data, flat_contents):
+    #     """Reconstructs a variable from the aux_data and the leaves.
+    #
+    #     Args:
+    #       aux_data:
+    #       flat_contents:
+    #
+    #     Returns:
+    #       The variable.
+    #     """
+    #     return cls(*flat_contents)
+
+    # def clone(self) -> 'Variable':
+    #     """Clone the variable. """
+    #     r = type(self)(jnp.array(self.value, copy=True), batch_axis=self.batch_axis)
+    #     return r
+
+    # def __eq__(self, other):
+    #     """Override State's __eq__ to use BaseArray behavior for element-wise comparison."""
+    #     from brainpy._src.math.ndarray import _check_input_array, _return
+    #     return _return(self.value == _check_input_array(other))
+    #
+    # def __ne__(self, other):
+    #     """Override State's __ne__ to use BaseArray behavior for element-wise comparison."""
+    #     from brainpy._src.math.ndarray import _check_input_array, _return
+    #     return _return(self.value != _check_input_array(other))
+    #
+    # def __lt__(self, other):
+    #     """Override State's __lt__ to use BaseArray behavior for element-wise comparison."""
+    #     from brainpy._src.math.ndarray import _check_input_array, _return
+    #     return _return(self.value < _check_input_array(other))
+    #
+    # def __le__(self, other):
+    #     """Override State's __le__ to use BaseArray behavior for element-wise comparison."""
+    #     from brainpy._src.math.ndarray import _check_input_array, _return
+    #     return _return(self.value <= _check_input_array(other))
+    #
+    # def __gt__(self, other):
+    #     """Override State's __gt__ to use BaseArray behavior for element-wise comparison."""
+    #     from brainpy._src.math.ndarray import _check_input_array, _return
+    #     return _return(self.value > _check_input_array(other))
+    #
+    # def __ge__(self, other):
+    #     """Override State's __ge__ to use BaseArray behavior for element-wise comparison."""
+    #     from brainpy._src.math.ndarray import _check_input_array, _return
+    #     return _return(self.value >= _check_input_array(other))
 
 
 def _get_dtype(v):

brainpy/_src/train/back_propagation.py

@@ -230,6 +230,8 @@ def fit(
              Please set batch size in your dataset.
 
         """
+
+        share.save(fit=True)
         if shared_args is None:
             shared_args = dict()
         shared_args['fit'] = shared_args.get('fit', True)

examples/dynamics_simulation/stdp.py

@@ -20,16 +20,20 @@ def __init__(self, num_poisson, num_lif=1, g_max=0.01):
 
         # neuron groups
         self.noise = bp.dyn.PoissonGroup(num_poisson, freqs=15.)
-        self.group = bp.dyn.Lif(num_lif, V_reset=-60., V_rest=-74, V_th=-54, tau=10.,
-                                V_initializer=bp.init.Normal(-60., 1.))
+        self.group = bp.dyn.Lif(
+            num_lif, V_reset=-60., V_rest=-74, V_th=-54, tau=10.,
+            V_initializer=bp.init.Normal(-60., 1.)
+        )
 
         # synapses
         syn = bp.dyn.Expon.desc(num_lif, tau=5.)
         out = bp.dyn.COBA.desc(E=0.)
-        comm = bp.dnn.AllToAll(num_poisson, num_lif, bp.init.Uniform(0., g_max))
-        self.syn = bp.dyn.STDP_Song2000(self.noise, None, syn, comm, out, self.group,
-                                        tau_s=20, tau_t=20, W_max=g_max, W_min=0.,
-                                        A1=0.01 * g_max, A2=0.0105 * g_max)
+        comm = bp.dnn.AllToAll(num_poisson, num_lif, bp.init.Uniform(0., g_max), mode=bm.TrainingMode())
+        self.syn = bp.dyn.STDP_Song2000(
+            self.noise, None, syn, comm, out, self.group,
+            tau_s=20, tau_t=20, W_max=g_max, W_min=0.,
+            A1=0.01 * g_max, A2=0.0105 * g_max
+        )
 
     def update(self, *args, **kwargs):
         self.noise()

examples/dynamics_training/echo_state_network.py

@@ -4,6 +4,7 @@
 import brainpy.math as bm
 
 bm.set_environment(bm.batching_mode)
+bp.share.save(fit=True)
 
 
 class ESN(bp.DynamicalSystem):

examples/dynamics_training/reservoir-mnist.py

@@ -8,8 +8,8 @@
 import brainpy as bp
 import brainpy.math as bm
 
-traindata = bd.vision.MNIST(root='D:/data', split='train')
-testdata = bd.vision.MNIST(root='D:/data', split='test')
+traindata = bd.vision.MNIST(root='./data', split='train', download=True)
+testdata = bd.vision.MNIST(root='./data', split='test', download=True)
 
 
 def offline_train(num_hidden=2000, num_in=28, num_out=10):
@@ -18,7 +18,7 @@ def offline_train(num_hidden=2000, num_in=28, num_out=10):
     x_train = x_train.reshape(-1, x_train.shape[-1])
     y_train = bm.one_hot(jnp.repeat(traindata.targets, x_train.shape[1]), 10, dtype=bm.float_)
 
-    reservoir = bp.layers.Reservoir(
+    reservoir = bp.dyn.Reservoir(
         num_in,
         num_hidden,
         Win_initializer=bp.init.Uniform(-0.6, 0.6),

examples/operator_customization/event_ell.py

@@ -2,6 +2,6 @@ numpy
 brainstate>=0.1.6
 brainunit
 brainevent>=0.0.4
-braintools>=0.0.7
+braintools>=0.0.9
 jax
 tqdm

requirements.txt

@@ -33,7 +33,7 @@
     packages=packages,
     python_requires='>=3.10',
     install_requires=[
-        'numpy>=1.15', 'jax', 'tqdm', 'brainstate>=0.1.6', 'brainunit', 'brainevent>=0.0.4', 'braintools>=0.0.7'
+        'numpy>=1.15', 'jax', 'tqdm', 'brainstate>=0.1.6', 'brainunit', 'brainevent>=0.0.4', 'braintools>=0.0.9'
     ],
     url='https://github.com/brainpy/BrainPy',
     project_urls={