brainpy
Leak problem of the DSRunner
Hello brainpy team: I found a leak problem in your sample code:
with jax.checking_leaks():
model = bp.dyn.HH(3, gNa=bp.init.Uniform(min_val=100, max_val=140))
inputs = np.ones(int(100./ bm.dt)) * 6. # 100 ms
runner = bp.DSRunner(model, monitors=['V'])
runner.run(inputs=inputs)
it caused:
Exception: Leaked trace DynamicJaxprTrace. Leaked tracer(s):
Traced<ShapedArray(uint32[2])>with<DynamicJaxprTrace>
<DynamicJaxprTracer 2107853421152> is referred to by <RandomState 2106157584272>
<RandomState 2106157584272> is referred to by brainpy._src.math.random.DEFAULT
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
I noticed that you also mentioned the memory leak problem of DSRunner, and I use it in my project. My question is that would it make the exit code:-1073741819 (0xC0000005) or taichi error message happen?
When i used the env In my project: ● python == 3.11 ● brainpy == 2.6.0.post20240618[cpu] ● jax and jaxlib == 0.4.31[cpu]
It always reported exit code:-1073741819 (0xC0000005). But when i use the env: ● python == 3.11 ● brainpy == 2.6.0.post20241205[cpu] ● jax and jaxlib == 0.4.36[cpu]
There are three kind of situations: ● First, the process runs normally ● Second, just report exit code:-1073741819 (0xC0000005) ● Third, taichi reports:
[W 12/09/24 14:46:06.384 52476] [taichi_core_impl.cpp:ti_set_last_error@233] C-API error: (argument out of range) args[1].type
[W 12/09/24 14:46:06.385 95120] [taichi_core_impl.cpp:ti_set_last_error@233] C-API error: (invalid state) Access violation - no RTTI data!
By the way, my friend can run this process normally. I used their env but it did not work. Is it my computer broken?
A easy way to fix the leaky problem about random.DEFAULT is reseting the random seed.
For example:
with jax.checking_leaks():
bm.random.seed()
model = bp.dyn.HH(3, gNa=bp.init.Uniform(min_val=100, max_val=140))
inputs = np.ones(int(100./ bm.dt)) * 6. # 100 ms
runner = bp.DSRunner(model, monitors=['V'])
runner.run(inputs=inputs)
I used the latest versions of brainunit and brainstate to test 5 times in different environments, and they all succeeded. As for random seed, it did not work.