Memory leak in jx.integrate or jx.Cell

[jnsbck]

Hi,

while running a bunch of sims in parallel, I noticed that the memory consumption went through the roof. When I looked into it I could reduce it to just running jx.integrate in a while loop, see snippet below. Do I need to clear some cache manually somehow or is memory leaking somewhere? If the latter is the case than I think a fix should be high priority.

from jax import config
from jax import jit

config.update("jax_enable_x64", True)
config.update("jax_platform_name", "cpu")

import os
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = ".1"

import numpy as np
import jaxley as jx
from jaxley.channels import HH
import psutil

cell = jx.read_swc(f"morphologies/20161028_1.swc", nseg=4, min_radius=5.0)
cell.insert(HH())

cell.delete_recordings()
cell.branch(1).comp(0).record()

 # @jit # jitting seems to plug the leakage
def simulate():
    return jx.integrate(cell, t_max=1.0)

if __name__ == "__main__":
    while True:
        voltage = simulate()
        print(f"Memory usage: {psutil.virtual_memory().percent}%")

things I already tried:

• resetting recordings before every integration.
• using different hardware
• the rate at which it leaks seems to be independent of number of t_steps / t_max

Lemme know if there is any other tests I should run or if you cannot reproduce this.

EDIT:
@kyralianaka found out that jitting the function fixes the issue. My hypothesis is that each call to integrate instantiates some new objects that don't get destroyed for some reason. In the jitted version, this does not happen since they are only instantiated once.

[michaeldeistler]

Interesting, thanks for reporting. I do not know why exactly this happens tbh.