Skip to content

Commit

Permalink
[dyn] add clear_input in the step_run function (#601)
Browse files Browse the repository at this point in the history
Usually, we use `step_run` with `brainpy.math.for_loop`. This function does not call `brainpy.clear_input`, which may cause problems of input accumulation when using old APIs in ``brainpy.neurons`` module. Therefore, we should call ``brainpy.clear_input`` after each ``update`` function.
  • Loading branch information
chaoming0625 authored Jan 26, 2024
1 parent bc5aa72 commit 8c57f66
Showing 1 changed file with 11 additions and 2 deletions.
13 changes: 11 additions & 2 deletions brainpy/_src/dynsys.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@
IonChaDyn = None
SLICE_VARS = 'slice_vars'
the_top_layer_reset_state = True
clear_input = None
reset_state = None


def not_implemented(fun):
Expand Down Expand Up @@ -146,7 +148,9 @@ def reset(self, *args, **kwargs):
See https://brainpy.readthedocs.io/en/latest/tutorial_toolbox/state_resetting.html for details.
"""
from brainpy._src.helpers import reset_state
global reset_state
if reset_state is None:
from brainpy._src.helpers import reset_state
reset_state(self, *args, **kwargs)

@not_implemented
Expand Down Expand Up @@ -178,8 +182,13 @@ def step_run(self, i, *args, **kwargs):
Returns:
out: The update function returns.
"""
global clear_input
if clear_input is None:
from brainpy._src.helpers import clear_input
share.save(i=i, t=i * bm.dt)
return self.update(*args, **kwargs)
out = self.update(*args, **kwargs)
clear_input(self)
return out

@bm.cls_jit(inline=True)
def jit_step_run(self, i, *args, **kwargs):
Expand Down

0 comments on commit 8c57f66

Please sign in to comment.