Leak problem of the DSRunner

[nancheng233]

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?

[chaoming0625]

Thanks for report. Can you try brainstate?

[chaoming0625]

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)

[nancheng233]

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.

[chaoming0625]

Thanks for your feedback.