Easy-Transformer icon indicating copy to clipboard operation
Easy-Transformer copied to clipboard

[Proposal] change FactoredMatric.svd() so it doesn't prevent all instances of FactoredMatrix from being garbage collected

Open manulari opened this issue 1 year ago • 1 comments

Proposal

FactoredMatrix.svd is annotated with @lru_cache. This prevents instances of FactoredMatrix from ever being garbage collected, as explained here: https://rednafi.com/python/lru_cache_on_methods/

This post also gives a solution, which is to remove the annotation on the method definition, and instead do

def __init__(self, A, B):
    ...
    self.svd = lru_cache()(self._svd)
    ...

def _svd(self):
    ...

Additional context

There is more discussion of the issue in this stackoverflow post. The python docs also have a faq entry

The solution above still creates a cyclical reference, so requires cycle detection by the garbage collector.

Alternative solution

In this specific case it probably makes more sense to use the @cached_property-decorator from functools instead, as it is in the standard library and designed to be used per instance. It avoids a cyclic reference.

@cached_property
def svd(self):
    ...

Or if for backwards compatibility one wants svd to remain a method and not a property, then one could do:

@cached_property
def _svd(self):
    ...
def svd(self):
    return self._svd

Note that svd being a method is a bit inconsistent anyways, as U, S, Vh are properties. Also, eigenvalues could be a @cached_property as well.

Checklist

  • [x] I have checked that there is no similar issue in the repo (required)

manulari avatar Nov 25 '24 11:11 manulari

Thank you for bringing this up. We can probably work this into the next major release of TransformerLens.

bryce13950 avatar Nov 25 '24 21:11 bryce13950