Typing error using LanguageModel.trace() as a context manager

[chanind]

The following code works, but pyright/pylance gives a typing error:

from nnsight import LanguageModel

model = LanguageModel('google/gemma-2-2b')

with model.trace('hi'):
    acts = model.model.layers[0].output.save()

The errors are the following:

Object of type "InterventionProxy" cannot be used with "with" because it does not implement __enter__
  Attribute "__enter__" is unknown

Object of type "InterventionProxy" cannot be used with "with" because it does not implement __exit__

It looks like this because LanguageModel.__new__() is typed as returning either a LanguageModel or an Envoy. Envoy.__call__() returns an InterventionProxy, which is not a contextmanager. Since the type checker doesn't know if the model is a LanguageModel or an Envoy, it has to assume that model might be an Envoy and thus can't use with model.trace() as a contextmanager. It seems like __new__() should be typed to just return a LanguageModel, no?

[JadenFiotto-Kaufman]

@chanind Would you mind trying this on the 0.4 branch and let me know if it still complains? pip install --upgrade git+https://github.com/ndif-team/[email protected] If so I'll take a closer look.

In general though, I am doing very weird things with typing... this is because I have a lot of objects (like NNsight, Envoy, and Proxies) that represent other types of objects (Huggingface models, torch modules, and tensors respectively) but aren't themselves those objects. So I would like you IDE to still be able to complete/hint at methods/attributes on their "wrapper" objects. If you have any thoughts on this I'd love to hear them!