nnsight
nnsight copied to clipboard
Calling `next` after `all` should raise a proper error
Disclaimer: I was not surprised that this fails as it doesn't make sense to use all and next on the same module
from nnsight import LanguageModel
model = LanguageModel("gpt2")
with model.generate(prompt, max_new_tokens=10):
model.transformer.h.all()
model.transformer.h.next()
Fails with TypeError: unsupported operand type(s) for +: 'slice' and 'int'
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
File /workspace/clement/repos/representation-structure-comparison/junk.py:4
2 from nnsight import LanguageModel
3 model = LanguageModel("gpt2")
----> 4 with model.generate(prompt, max_new_tokens=10):
5 model.transformer.h.all()
6 model.transformer.h.next()
File /workspace/clement/repos/representation-structure-comparison/.env/lib/python3.10/site-packages/nnsight/intervention/contexts/interleaving.py:96, in InterleavingTracer.__exit__(self, exc_type, exc_val, exc_tb)
92 self.invoker.__exit__(None, None, None)
94 self._model._envoy._reset()
---> 96 super().__exit__(exc_type, exc_val, exc_tb)
File /workspace/clement/repos/representation-structure-comparison/.env/lib/python3.10/site-packages/nnsight/tracing/contexts/tracer.py:25, in Tracer.__exit__(self, exc_type, exc_val, exc_tb)
21 from .globals import GlobalTracingContext
23 GlobalTracingContext.try_deregister(self)
---> 25 return super().__exit__(exc_type, exc_val, exc_tb)
File /workspace/clement/repos/representation-structure-comparison/.env/lib/python3.10/site-packages/nnsight/tracing/contexts/base.py:72, in Context.__exit__(self, exc_type, exc_val, exc_tb)
69 graph = self.graph.stack.pop()
71 if isinstance(exc_val, BaseException):
---> 72 raise exc_val
74 self.add(graph.stack[-1], graph, *self.args, **self.kwargs)
76 if self.backend is not None:
File /workspace/clement/repos/representation-structure-comparison/junk.py:6
4 with model.generate(prompt, max_new_tokens=10):
5 model.transformer.h.all()
----> 6 model.transformer.h.next()
File /workspace/clement/repos/representation-structure-comparison/.env/lib/python3.10/site-packages/nnsight/intervention/envoy.py:251, in Envoy.next(self, increment)
241 def next(self, increment: int = 1) -> Envoy:
242 """By default, this modules inputs and outputs only refer to the first time its called. Use `.next()`to select which iteration .input an .output refer to.
243
244 Args:
(...)
248 Envoy: Self.
249 """
--> 251 return self.iter[self._iteration_stack[-1] + increment].__enter__()
TypeError: unsupported operand type(s) for +: 'slice' and 'int'
For newcomer, it'd be probably better if this raise a proper error message indicating that next should not be used on a module affected by all