JaxTyping introduces reference cycles that would otherwise not exist

[ojw28]

Particularly when building large/complicated systems, it can end up being highly preferable to delete Python objects by decrementing their reference counts to 0, as opposed to relying on the Python Garbage Collector (which is needed to delete unreachable objects that are either in or are referenced by a reference cycle).

For this reason, it would be really nice if decorating code with JaxTyping could avoid introducing any reference cycles where reference cycles do not already exist in the code being decorated. This is not currently the case. There are at least two places I know of where JaxTyping introduces reference cycles even if the user code does not contain any:

1. https://github.com/patrick-kidger/jaxtyping/blob/v0.2.34/jaxtyping/_decorator.py#L509 (wrapped_fn refering back to itself)
2. https://github.com/patrick-kidger/jaxtyping/blob/v0.2.34/jaxtyping/_decorator.py#L706 (there's a fn <-> scope reference cycle here)

Do you think avoiding reference cycles is a nice property to aim for, and do you think it should be achievable?

I think it should be possible to write a test helper that would let you do something like:

with assert_no_reference_cycles():
  # some code

and will tell you what the reference cycles are if it fails. Let me know if that would be helpful, and I can look into sending a pull request.

[patrick-kidger]

Thanks for the report! I agree, being able to avoid reference cycles would be nice.

I suspect the first one could be broken by using a weak reference. We really do need to be able to get wrapped_fn to check if it has been mutated with this attribute since it was created... but also there's no reason to create a cycle for that, if we're inside this function then it definitely still exists.

(Well, two caveats here:

1. It would theoretically be possible for that to no longer be true if we were to create a new function using an old types.CodeType object. Believe it or not I actually do have a use-case for that... but I also think that's weird enough that that's probably not jaxtyping's job to handle.
2. I'm not sure if changing this would affect anything in practice -- I'm guessing decorated functions will largely be global functions, which will continue to live until interpreter shutdown anyway.
   )

Meanwhile I suspect the second one could be manually broken via something like fn.__globals__ = {} (untested).

If you're willing to submit a PR + test then I'd be very happy to take that :)

[ojw28]

Regarding the second one:

fn.__globals__ = {}

causes a lot of tests to fail. I was able to "fix" it by adding a del after retrieving the function instead:

fn = scope[name]
del scope[name]

That said, I don't fully understand what this code is doing, so I don't know whether I'm breaking something else :). Does the del make sense as a thing to do?

[patrick-kidger]

Yup I think that del makes sense to me!

[ojw28]

Second change: #260

[patrick-kidger]

Okay, I think this is done! :) Are there any other examples you're bumping into or shall we close this?

[ojw28]

There aren't any other examples that I'm bumping into. If I find any, I'll send some more pull requests, but no need to keep this issue open just in case. Thanks for the prompt reviews and helping to get these changes landed!