Skip to content

Refactor STDP and event ops for new brainevent binary backend #817

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chaoming0625 merged 1 commit into from
Mar 11, 2026
Merged

Refactor STDP and event ops for new brainevent binary backend #817

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 2 additions & 8 deletions brainpy/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
# limitations under the License.
# ==============================================================================

__version__ = "2.7.6"
__version__ = "2.7.7"
__version_info__ = tuple(map(int, __version__.split(".")))

from brainpy import _errors as errors
Expand Down Expand Up @@ -133,20 +133,14 @@
synouts, # synaptic output
synplast, # synaptic plasticity
)
from brainpy.math.object_transform.base import (
Base as Base,
)
from brainpy.math.object_transform.base import Base as Base

from brainpy.math.object_transform.collectors import (
ArrayCollector as ArrayCollector,
Collector as Collector,
)

from brainpy.deprecations import deprecation_getattr

optimizers = optim


# New package
from brainpy import state

24 changes: 14 additions & 10 deletions brainpy/dnn/linear.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,12 @@
import jax
import jax.numpy as jnp
import numpy as np
from brainevent import csr_on_pre, csr2csc_on_post
from brainevent import dense_on_pre, dense_on_post
from brainevent import (
update_csr_on_binary_pre,
update_csr_on_binary_post,
update_dense_on_binary_pre,
update_dense_on_binary_post,
)

from brainpy import connect, initialize as init
from brainpy import math as bm
Expand Down Expand Up @@ -226,11 +230,11 @@ def stdp_update(
if on_pre is not None:
spike = on_pre['spike']
trace = on_pre['trace']
self.W.value = dense_on_pre(self.W.value, spike, trace, w_min, w_max)
self.W.value = update_dense_on_binary_pre(self.W.value, spike, trace, w_min, w_max)
if on_post is not None:
spike = on_post['spike']
trace = on_post['trace']
self.W.value = dense_on_post(self.W.value, trace, spike, w_min, w_max)
self.W.value = update_dense_on_binary_post(self.W.value, trace, spike, w_min, w_max)


Linear = Dense
Expand Down Expand Up @@ -321,11 +325,11 @@ def stdp_update(
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)
self.weight.value = update_dense_on_binary_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, trace, spike, w_min, w_max)
self.weight.value = update_dense_on_binary_post(self.weight.value, trace, spike, w_min, w_max)


class OneToOne(Layer, SupportSTDP):
Expand Down Expand Up @@ -449,11 +453,11 @@ def stdp_update(
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)
self.weight.value = update_dense_on_binary_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, trace, spike, w_min, w_max)
self.weight.value = update_dense_on_binary_post(self.weight.value, trace, spike, w_min, w_max)


class _CSRLayer(Layer, SupportSTDP):
Expand Down Expand Up @@ -500,7 +504,7 @@ def stdp_update(
if on_pre is not None: # update on presynaptic spike
spike = on_pre['spike']
trace = on_pre['trace']
self.weight.value = csr_on_pre(
self.weight.value = update_csr_on_binary_pre(
self.weight.value, self.indices, self.indptr, spike, trace, w_min, w_max,
shape=(spike.shape[0], trace.shape[0]),
)
Expand All @@ -512,7 +516,7 @@ def stdp_update(
)
spike = on_post['spike']
trace = on_post['trace']
self.weight.value = csr2csc_on_post(
self.weight.value = update_csr_on_binary_post(
self.weight.value, self._pre_ids, self._post_indptr,
self.w_indices, trace, spike, w_min, w_max,
shape=(trace.shape[0], spike.shape[0]),
Expand Down
4 changes: 3 additions & 1 deletion brainpy/initialize/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# -*- coding: utf-8 -*-
# Copyright 2025 BrainX Ecosystem Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand All @@ -13,6 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

# -*- coding: utf-8 -*-

import abc

__all__ = [
Expand Down
3 changes: 1 addition & 2 deletions brainpy/math/compat_numpy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@
'lcm', 'gcd', 'arccos', 'arccosh', 'arcsin', 'arcsinh', 'arctan',
'arctan2', 'arctanh', 'cos', 'cosh', 'sin', 'sinc', 'sinh', 'tan',
'tanh', 'deg2rad', 'hypot', 'rad2deg', 'degrees', 'radians', 'round',
'around', 'round_', 'rint', 'floor', 'ceil', 'trunc', 'fix', 'prod',
'around', 'round_', 'rint', 'floor', 'ceil', 'trunc', 'prod',
'sum', 'diff', 'median', 'nancumprod', 'nancumsum', 'nanprod', 'nansum',
'cumprod', 'cumsum', 'ediff1d', 'cross', 'isfinite', 'isinf',
'isnan', 'signbit', 'copysign', 'nextafter', 'ldexp', 'frexp', 'convolve',
Expand Down Expand Up @@ -397,7 +397,6 @@ def msort(a):
floor = _compatible_with_brainpy_array(jnp.floor)
ceil = _compatible_with_brainpy_array(jnp.ceil)
trunc = _compatible_with_brainpy_array(jnp.trunc)
fix = _compatible_with_brainpy_array(jnp.fix)
prod = _compatible_with_brainpy_array(jnp.prod)

sum = _compatible_with_brainpy_array(jnp.sum)
Expand Down
2 changes: 1 addition & 1 deletion brainpy/math/event/csr_matmat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ def csrmm(
if isinstance(matrix, Array):
matrix = matrix.value

matrix = brainevent.EventArray(matrix)
matrix = brainevent.BinaryArray(matrix)
csr = brainevent.CSR((data, indices, indptr), shape=shape)
if transpose:
return matrix @ csr
Expand Down
2 changes: 1 addition & 1 deletion brainpy/math/event/csr_matvec.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ def csrmv(
if isinstance(events, Array):
events = events.value

events = brainevent.EventArray(events)
events = brainevent.BinaryArray(events)
csr = brainevent.CSR((data, indices, indptr), shape=shape)
if transpose:
return events @ csr
Expand Down
8 changes: 4 additions & 4 deletions brainpy/math/jitconn/event_matvec.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,8 @@ def event_mv_prob_homo(
if isinstance(weight, Array):
weight = weight.value

events = brainevent.EventArray(events)
csr = brainevent.JITCHomoR((weight, conn_prob, seed), shape=shape, corder=outdim_parallel)
events = brainevent.BinaryArray(events)
csr = brainevent.JITCScalarR((weight, conn_prob, seed), shape=shape, corder=outdim_parallel)
if transpose:
return events @ csr
else:
Expand All @@ -75,7 +75,7 @@ def event_mv_prob_uniform(
seed = np.random.randint(0, 1000000000)
if isinstance(events, Array):
events = events.value
events = brainevent.EventArray(events)
events = brainevent.BinaryArray(events)
if isinstance(w_low, Array):
w_low = w_low.value
if isinstance(w_high, Array):
Expand Down Expand Up @@ -106,7 +106,7 @@ def event_mv_prob_normal(
seed = np.random.randint(0, 1000000000)
if isinstance(events, Array):
events = events.value
events = brainevent.EventArray(events)
events = brainevent.BinaryArray(events)
if isinstance(w_mu, Array):
w_mu = w_mu.value
if isinstance(w_sigma, Array):
Expand Down
4 changes: 2 additions & 2 deletions brainpy/math/jitconn/matvec.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ def mv_prob_homo(
if isinstance(weight, Array):
weight = weight.value

csr = brainevent.JITCHomoR((weight, conn_prob, seed), shape=shape, corder=outdim_parallel)
csr = brainevent.JITCScalarR((weight, conn_prob, seed), shape=shape, corder=outdim_parallel)
if transpose:
return vector @ csr
else:
Expand Down Expand Up @@ -290,7 +290,7 @@ def get_homo_weight_matrix(
"""
if seed is None:
seed = np.random.randint(0, 1000000000)
csr = brainevent.JITCHomoR((weight, conn_prob, seed), shape=shape, corder=outdim_parallel)
csr = brainevent.JITCScalarR((weight, conn_prob, seed), shape=shape, corder=outdim_parallel)
if transpose:
csr = csr.T
return csr.todense()
Expand Down
4 changes: 2 additions & 2 deletions brainpy/state/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# `brainpy.state` - State-based Brain Dynamics Programming
# `brainpy.state`

## Overview

Expand Down Expand Up @@ -37,7 +37,7 @@ pip install brainpy -U
For development or to install the state module separately:

```bash
pip install brainpy_state -U
pip install brainpy.state -U
Copy link

@sourcery-ai sourcery-ai bot Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

issue (bug_risk): The pip install brainpy.state command may be incorrect; consider confirming the actual package name on PyPI.

Because this is user-facing install guidance, an incorrect distribution name will lead to failed installs. Please verify the actual PyPI package name (e.g., whether it should be brainpy_state or another valid name) and update the command accordingly.

```

## Usage
Expand Down
5 changes: 3 additions & 2 deletions brainpy/state/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,5 +16,6 @@
from brainpy_state import *
from brainpy_state import __all__



if __name__ == '__main__':
print(LIF)
print(__all__)
3 changes: 1 addition & 2 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ classifiers = [
"Operating System :: OS Independent",
"Programming Language :: Python",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
Expand All @@ -41,7 +40,7 @@ dependencies = [
"tqdm",
"brainstate>=0.2.7",
"brainunit",
"brainevent>=0.0.4",
"brainevent>=0.0.7",
"braintools>=0.0.9",
'brainpy_state>=0.0.3',
]
Expand Down
2 changes: 1 addition & 1 deletion requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
numpy
brainunit
brainevent>=0.0.4
brainevent>=0.0.7
braintools>=0.1.0
brainstate>=0.2.7
brainpy_state>=0.0.2
Expand Down
Loading