[Proposal] Type hint support for self.model in ActivationCache

[Ja1Zhou]

Proposal

I propose to do a simple refractoring of the class HookedTransformer to support type hints for the model attribute in the ActivationCache class.

Motivation

I am frastrated by the lack of type hinting for self.model in the ActivationCache class. I think the current state is that we want to avoid circular imports, but I think we can do better.

Pitch

Create a class, say HookedTransformerMixin, that holds all methods without accessing to the definition of ActivationCache. Currently only the function run_with_cache access the ActivationCache class.

In transformer_lens/ActivationCache.py, we could import HookedTransformerMixin to provide type hint. In transformer_lens/HookedTransformer.py, HookedTransformer could inherit HookedTransformerMixin and import from ActivationCache.

Checklist

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

[bryce13950]

There is a low priority task to separate the HookedRootModule into it's own file, which could then be used as described, as well as making it easier for people to boot TransformerLens while receiving an instance of HookedRootModule based on config without necessarily needing to know which HookedModule needs to be manually instantiated. That is currently on the docket for 4.0, but we can probably work it into 3.0 in order to address this as well, providing that it doesn't turn into a large project.