Conditional Interventions#

Interventions can also be made conditional.

Inside the tracing context, we can specify a conditional context:

with tracer.cond(Boolean):
  pass

This context will only execute its contained interventions if the specified condition is met. Let’s try an example!

[3]:

import torch
from nnsight import LanguageModel

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

with model.trace("The Eiffel Tower is in the city of") as tracer:

  rand_int = torch.randint(low=-10, high=10, size=(1,))

  with tracer.cond(rand_int % 2 == 0):
    tracer.log("Random Integer ", rand_int, " is Even")

  with tracer.cond(rand_int % 2 == 1):
    tracer.log("Random Integer ", rand_int, " is Odd")

Random Integer  tensor([4])  is Even

In the example above, we have two conditional contexts with mutually exclusive conditions, just like a usual If-Else statement.

Conditional contexts can also be nested, if we want our interventions to depend on more than one condition at a time.

[4]:

with model.trace("The Eiffel Tower is in the city of") as tracer:

  non_rand_int = 8

  with tracer.cond(non_rand_int > 0):
    with tracer.cond(non_rand_int % 2 == 0):
      tracer.log("Rand Int ", non_rand_int, " is Positive and Even")

Rand Int  8  is Positive and Even

nnsight 0.4 introduces support for native Python if statements within the tracing context! Simply create an if statement within a trace, and it should perform as tracer.cond().

[6]:

with model.trace("The Eiffel Tower is in the city of") as tracer:

  rand_int = torch.randint(low=-10, high=10, size=(1,))

  # Since this if statement is inside the tracing context the if will
  # create a conditional context and will only execute the intervention
  # if this condition is met
  if rand_int % 2 == 0:
    tracer.log("Random Integer ", rand_int, " is Even")

  if rand_int % 2 == 1:
    tracer.log("Random Integer ", rand_int, " is Odd")

Random Integer  tensor([0])  is Even

Considerations#

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

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

[ ]:

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

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