leak reports on hello_word function when using jax.jit

[Char-Aznable]

I'm not sure if this is a nanobind issue but just posting in case someone has seen this. Basically, I was following the jax documentation here to make custom c++/cuda calls in jax but using nanobind instead of pybind11 as in the doc. The code works fine but I see a whole bunch of leak warnings from nanobind only when the functions are called by jax.jit(myFunction)(args). No leak warning if I just call myFunction(args), where myFunction in turn calls the python bindings. Has anyone seen similar things? How do I make sense of it?

As a negative control, I make a hello world function that boils down to:

NB_MODULE(XformLigand, m) {
  m.def("make_MetaDataHelloWord", []() { return; });
}

and I can see nanobind report leaks on this function in jax.jit mode

[wjakob]

Nanobind is pretty noisy about leaks—with pybind11, those leaks probably also occurred, but you simply did not know about it.

I suppose that @jax.jit is sort of a special function in that it traces code and then has to hold onto that traced code in case somebody later calls the function again with similar arguments. Perhaps in doing so, it keeps a reference to a nanobind function or an object, and doesn't clean it up.

Your Python function wrapper might isolate you from doing that since it doesn't reference a specific function but rather includes the code to create that reference at a later point.

Ultimately, if your goal is to interface with something like JAX, you can either push the developers of this tool to be more conscious of cleanup and garbage collection when the extension library shuts down. Or you can disable the nanobind warnings (see the docs). If you change nanobind gobal flags like that, then it would be good if you could compile it with NB_DOMAIN specified in the CMakeLists.txt

[Char-Aznable]

What confuses me is the report message leaked function "make_MetaDataHelloWord" that makes look like the function itself leaks but judging from what you said it's the calling context that leaks. Is there way to make better leak report e.g., showing the stack trace for the leak? or is there a way to identify which exact function is leaking?

[Char-Aznable]

A more common usage case where I have seen this kind of false positive report is using pytest to test nanobind functions. I have also done negative controls in the pytest cases but I often find those leak reports not very robust, i.e., only occur sporadically. Is this also expected? I suppose pytest has a much larger user base than jax so someone might know about this issue

[wjakob]

Python is no longer alive when nanobind reports these leaks, so the ability to acquire additional information is unfortunately too limited to do what you suggest. Have you tried raising this issue with Jax? Presumably they have some kind of cache that they don't clean up at shutdown time.