Clearing cache of traced/jitted functions

Hi all,

Is there a public way to clear/invalidate the traced functions? In order to a) retrace it (with a changed Python state) and b) clear up the memory

The use-case is that a function is sometimes (Python) state based and therefore one would need to retrace it, basically invalidating the old trace and removing it.

I am aware that there are often ways around such a design, but sometimes there are None: if you build a model that depends on Python variable and run this a couple of times. Manually keeping track of changes is one thing (this is up to the user, clearly) but we cannot simply check for re-tracing in general (since this would mean that we would also need to handle the re-tracing logic of Tensors).

How could we go about clearing the traced functions?

(I managed to “hack” around and “get something working” but accessing the private interface, so that’s probably not the way forward)