search

ndif-team / nnsight

The nnsight package enables interpreting and manipulating the internals of deep learned models.
https://nnsight.net/
MIT License
356 stars 34 forks source link

Conditional Context has a GraphBasedContext entry point #197

Closed AdamBelfki3 closed 3 weeks ago

AdamBelfki3 commented 3 weeks ago

Reverting to the previous API design for creating Conditional contexts, using any graph-based context.

Implementation Changes

ConditionalProtocol nodes don't have to strictly depend on Intervention Proxies anymore, as the condition input can be of any data type and will be evaluated with bool(). This means that you can have ConditionalProtocol nodes directly depend on other ConditionalProtocol nodes.

InterventionProxy is no longer a ContextManager for the moment.

Usage

with tiny_model.trace(tiny_input) as tracer:
        num = 5
        with tracer.cond(num > 0):
            tiny_model.layer1.output[:] = 1
            l1_out = tiny_model.layer1.output.save()