Iterative Interventions#

NNsight’s iterator context allows us to run an intervention loop at scale. It iteratively executes and updates a single intervention graph.

Use a session to define the Iterator context and pass in a sequence of items that you want to loop over at each iteration:

[7]:

import nnsight
from nnsight import LanguageModel

model = LanguageModel('openai-community/gpt2', device_map='auto')


with model.session() as session:

  with session.iter([0, 1, 2]) as item:
    # define intervention body here ...

    with model.trace("_"):
      # define interventions here ...
      pass

    with model.trace("_"):
      # define interventions here ...
      pass

The Iterator context extends all the nnsight graph-based functionalities, but also closely mimics the conventional for loop statement in Python, which allows it to support all kind of iterative operations with a use of as item syntax.

Beyond specifying iteration indices, you can also loop across an NNsight list object (nnsight.list()).

[8]:

with model.session() as session:

  li = nnsight.list() # an NNsight built-in list object
  [li.append([num]) for num in range(0, 3)] # adding [0], [1], [2] to the list
  li2 = nnsight.list().save()

  # You can create nested Iterator contexts
  with session.iter(li) as item:
    with session.iter(item) as item_2:
      li2.append(item_2)

print("\nList: ", li2)


List:  [0, 1, 2]

nnsight 0.4 introduces support for native Python for loops within a tracer context at scale!

NOTE: inline for loops (i.e., ``[x for x in <Proxy object>]``) are not currently supported.

[9]:

# New: Using Python for loops for iterative interventions
with model.session() as session:

    li = nnsight.list()
    [li.append([num]) for num in range(0, 3)]
    li2 = nnsight.list().save()

    # Using regular for loops
    for item in li:
        for item_2 in item: # for loops can be nested!
            li2.append(item_2)

print("\nList: ", li2)


List:  [0, 1, 2]

Considerations#

If you would like to turn off NNsight’s support of native for loops, you can apply the following changes to nnsight.CONFIG

This will not affect any of NNsight’s .iter() functionality.

[11]:

# Turn off support if/for statements within tracing context.
import nnsight

nnsight.CONFIG.APP.CONTROL_FLOW_HANDLING = False
nnsight.CONFIG.save()
On this page