envoy¶

envoy ¶

Envoy ¶

Envoy(module: Module, interleaver: Optional[Interleaver] = None, path: Optional[str] = 'model', rename: Optional[Dict[str, Union[str, List[str]]]] = None, envoys: Optional[Union[Type['Envoy'], Dict[Type[Module], Type['Envoy']]]] = None)

Bases: Batchable

A proxy class that wraps a PyTorch module to enable intervention during execution.

This class provides access to module inputs and outputs during forward passes, and allows for modification of these values through an interleaving mechanism. It serves as the primary interface for inspecting and modifying the behavior of neural network modules during execution.

ATTRIBUTE	DESCRIPTION
`path`	The module's location in the model hierarchy. Example: "model.encoder.layer1" indicates this module is the first layer of the encoder in the model. TYPE: `str`
`_module`	The underlying PyTorch module TYPE: `Module`
`_source`	Source code representation of the module TYPE: `Optional[EnvoySource]`
`interleaver`	Interleaver for managing execution flow TYPE: `Optional[Interleaver]`
`_default_mediators`	List of default mediators created with .edit TYPE: `List[List[str]]`
`_children`	List of child Envoys TYPE: `List[Envoy]`
`_alias`	Aliaser object for managing aliases TYPE: `Aliaser`

PARAMETER	DESCRIPTION
`module`	The PyTorch module to wrap TYPE: `Module`
`interleaver`	Optional interleaver for managing execution flow TYPE: `Optional[Interleaver]` DEFAULT: `None`
`path`	Optional path string representing the module's location in the model hierarchy TYPE: `Optional[str]` DEFAULT: `'model'`
`rename`	Optional dictionary mapping module names to alias names. Example: {"layer1": "first_layer", "layer2": "second_layer"} Example: {".model.layers": ".layers"} <-- Mounts .layers to the root model. Example: {".transformer": ["model", "mdl"]} <-- Allows access of .transformer as .model or .mdl TYPE: `Optional[Dict[str, Union[str, List[str]]]]` DEFAULT: `None`
`envoys`	Controls which Envoy class wraps descendant modules. Propagates down the envoy tree. - None (default): all descendants are wrapped with the base Envoy class. - A class: all descendants are wrapped with that class. - A dict whose values are `Envoy` subclasses. Keys may be: * A `torch.nn.Module` subclass — matches when the class appears in the descendant's MRO. Example: `{torch.nn.Linear: MyLinearEnvoy}`. * A string — matches when the descendant's envoy path ends with the key treated as a dotted suffix (component-wise). With a rename dict in play, each component also matches via single-component aliases — so `{"attn": MyAttnEnvoy}` matches a path ending in `self_attn` when the user passed `rename={"self_attn": "attn"}`. Type keys are tried first; string keys are a fallback. Descendants without a match fall back to the base Envoy class. Example: {torch.nn.Linear: MyLinearEnvoy, "self_attn": MyAttnEnvoy} TYPE: `Optional[Union[Type[Envoy], Dict]]` DEFAULT: `None`

path `instance-attribute` ¶

path = path

interleaver `instance-attribute` ¶

interleaver = interleaver if interleaver is not None else Interleaver()

interleaving `property` ¶

interleaving: bool

Check if the Envoy is currently nterleaving.

RETURNS	DESCRIPTION
`bool`	True if the Envoy is interleaving, False otherwise

source `property` ¶

source: SourceEnvoy

Get the source code representation of the module.

Lazily resolves to a :class:SourceEnvoy over the module's global :class:SourceAccessor. The accessor — and its per-call-site :class:OperationAccessor instances — are created once per module and shared across all Envoys / Interleavers / Mediators that touch it. This Envoy keeps a per-instance :class:SourceEnvoy so each Envoy has its own user-facing wrapper.

Examples:

>>> model = LanguageModel("gpt2", device_map='auto', dispatch=True)
>>> print(model.transformer.h[0].attn.source)
>>> with model.trace("Hello World"):
...     attn = model.transformer.h[0].attn.source.attention_interface_0.output.save()

RETURNS	DESCRIPTION
`A`	class:`SourceEnvoy` exposing operation-level access. TYPE: `SourceEnvoy`

iter `property` ¶

iter

device `property` ¶

device: Optional[device]

Get the device the module is on. Finds the first parameter and return its device.

devices `property` ¶

devices: Optional[set[device]]

Get the devices the module is on. Finds all parameters and return their devices.

getitem ¶

__getitem__(key: str) -> Envoy

Access a child Envoy by index for Module Lists.

PARAMETER	DESCRIPTION
`key`	The index of the child Envoy to retrieve TYPE: `str`

RETURNS	DESCRIPTION
`Envoy`	The child Envoy at the specified index

output ¶

output() -> Object

Get the output of the module's forward pass.

Examples:

>>> with model.trace("Hello World"):
...     attn = model.transformer.h[0].attn.output[0].save()

inputs ¶

inputs() -> Tuple[Tuple[Object], Dict[str, Object]]

Get the inputs to the module's forward pass.

RETURNS	DESCRIPTION
`Tuple[Tuple[Object], Dict[str, Object]]`	(args, kwargs) tuple of positional and keyword arguments.

Examples:

>>> with model.trace("Hello World"):
...     args, kwargs = model.transformer.h[0].attn.inputs

input ¶

input(value)

call ¶

__call__(*args, hook: bool = False, **kwargs)

trace ¶

trace(*args, trace: bool = True, fn: Optional[Callable] = None, tracer_cls: Type[InterleavingTracer] = InterleavingTracer, **kwargs)

Create a tracer for this module.

This method returns a tracer that can be used to capture and modify the execution of the module.

Examples:

>>> model = LanguageModel("gpt2", device_map='auto', dispatch=True)
>>> with model.trace("Hello World"):
...     model.transformer.h[0].attn.output[0][:] = 0

... output = model.output.save()

>>> print(output)

PARAMETER	DESCRIPTION
`*args`	Arguments to pass to the tracer DEFAULT: `()`
`trace`	If False, bypass tracing entirely — run the underlying module on the prepared input and return its output. Useful for one-shot forward passes that don't need intervention. Defaults to True. TYPE: `bool` DEFAULT: `True`
`**kwargs`	Keyword arguments to pass to the tracer DEFAULT: `{}`

RETURNS	DESCRIPTION
	An InterleavingTracer for this module, or — when `trace=False`
	— the module's output value directly.

scan ¶

scan(*args, **kwargs)

Just like .trace() but runs the model in fake tensor mode to validate operations and inspect tensor shapes.

This method returns a tracer that runs the model in fake tensor mode to validate operations and inspect tensor shapes without performing actual computation. This is useful for: - Validating that operations will work with given input shapes - Inspecting the shapes and types of tensors that would flow through the model - Debugging shape mismatches or other tensor-related issues.

Note this will not dispatch the model if not dispatched.

Examples:

>>> model = LanguageModel("gpt2", device_map='auto', dispatch=True)
>>> # Value error as the fake inputs and outputs have not been scanned in.
>>> print(model.transformer.h[0].mlp.output.shape)
>>> # Scan the model to validate operations and inspect shapes
>>> with model.scan("Hello World"):
...     # Access fake inputs/outputs to inspect shapes
...     attn_input = model.transformer.h[0].attn.input.save()
...     attn_output = model.transformer.h[0].attn.output[0].save()
>>> print(f"Attention input shape: {attn_input.shape}")
>>> print(f"Attention output shape: {attn_output.shape}")
>>> print(model.transformer.h[0].mlp.output.shape)

PARAMETER	DESCRIPTION
`*args`	Arguments to pass to the tracer DEFAULT: `()`
`**kwargs`	Keyword arguments to pass to the tracer DEFAULT: `{}`

RETURNS	DESCRIPTION
	A ScanningTracer for this module

edit ¶

edit(*, inplace: bool = False)

Create an editing tracer for this module. Allows for setting default interventions. This means this tracer won't execute the module, but will instead set default interventions that are applied on all future executions.

Edits can be cleared with Envoy.clear_edits().

Examples:

>>> model = LanguageModel("gpt2", device_map='auto', dispatch=True)
>>> # Now the first layer attention output will always be 0.
>>> with model.edit() as edited_model:
...     edited_model.transformer.h[0].attn.output[:] = 0

>>> with model.trace("Hello World"):
...     output = model.output.save()
>>> # The orignal model will have the default output.
>>> print(output)

>>> with edited_model.trace("Hello World"):
...     edited_output = edited_model.output.save()
>>> # The edited model will have the output after our intervention.
>>> print(edited_output)

PARAMETER	DESCRIPTION
`inplace`	Whether to edit in place. Defaults to False. TYPE: `bool` DEFAULT: `False`

RETURNS	DESCRIPTION
`EditingTracer`	An EditingTracer for this module

clear_edits ¶

clear_edits()

Clear all edits for this Envoy.

export_edits ¶

export_edits(name: str, export_dir: Optional[str] = None, variant: str = '__default__')

TODO

PARAMETER	DESCRIPTION
`name`	description TYPE: `str`
`export_dir`	description. Defaults to None. TYPE: `Optional[str]` DEFAULT: `None`
`variant`	description. Defaults to 'default'. TYPE: `str` DEFAULT: `'__default__'`

RAISES	DESCRIPTION
`ValueError`	description

import_edits ¶

import_edits(name: str, export_dir: Optional[str] = None, variant: str = '__default__')

TODO

PARAMETER	DESCRIPTION
`name`	description TYPE: `str`
`export_dir`	description. Defaults to None. TYPE: `Optional[str]` DEFAULT: `None`
`variant`	description. Defaults to 'default'. TYPE: `str` DEFAULT: `'__default__'`

session ¶

session(*args, tracer_cls: Type[Tracer] = Tracer, **kwargs)

all ¶

all()

next ¶

next(step: int = 1)

skip ¶

skip(replacement: Any)

Skips the execution of this module duting execution / interleaving. Behavior is the module will not be executed and will return a replacement value instead.

Examples:

>>> model = LanguageModel("gpt2", device_map='auto', dispatch=True)
>>> with model.trace("Hello World"):
...     # Skip the first layer and replace it with the input to the layer.
...     model.transformer.h[0].skip((model.transformer.h[0].input, None))
...     output = model.output.save()
>>> print(output)

PARAMETER	DESCRIPTION
`replacement`	The replacement value to replace the module's output with. TYPE: `Any`

to ¶

to(device: device)

Move the module to a specific device.

This method moves the underlying PyTorch module to the specified device.

PARAMETER	DESCRIPTION
`device`	The device to move the module to TYPE: `device`

RETURNS	DESCRIPTION
	Self, for method chaining

cpu ¶

cpu(*args, **kwargs)

Move the module to the CPU.

cuda ¶

cuda(*args, **kwargs)

Move the module to the GPU.

modules ¶

modules(include_fn: Callable[[Envoy], bool] = None, names: bool = False) -> List[Envoy]

Get all modules in the Envoy tree.

This method returns all Envoys in the tree, optionally filtered by an inclusion function.

PARAMETER	DESCRIPTION
`include_fn`	Optional function to filter modules TYPE: `Callable[[Envoy], bool]` DEFAULT: `None`
`names`	Whether to include module names in the result TYPE: `bool` DEFAULT: `False`

RETURNS	DESCRIPTION
`List[Envoy]`	A list of Envoys or (name, Envoy) tuples

named_modules ¶

named_modules(*args, **kwargs) -> List[Tuple[str, Envoy]]

Returns all Envoys in the Envoy tree along with their name/module_path.

This is a convenience method that calls modules() with names=True.

PARAMETER	DESCRIPTION
`include_fn`	Optional function to be ran against all Envoys to check if they should be included in the final collection of Envoys. Defaults to None. TYPE: `Callable`
`args, *kwargs`	Additional arguments to pass to modules()

RETURNS	DESCRIPTION
`List[Tuple[str, Envoy]]`	List[Tuple[str, Envoy]]: Included Envoys and their names/module_paths.

get ¶

get(path: str) -> Object

Gets the Envoy/Proxy via its path.

e.x: model = nnsight.LanguageModel("openai-community/gpt2")

module = model.get('transformer.h.0.mlp')

with model.trace("Hello"):
    value = model.get('transformer.h.0.mlp.output').save()

PARAMETER	DESCRIPTION
`path`	'.' separated path. TYPE: `str`

RETURNS	DESCRIPTION
`Object`	Union[Envoy, InterventionProxyType]: Fetched Envoy/Proxy

interleave ¶

interleave(fn: Union[Callable, str], *args, **kwargs)

len ¶

__len__()

Get the length of the Envoy.

iter ¶

__iter__()

Iterate over the Envoy.

str ¶

__str__()

String representation of the Envoy.

RETURNS	DESCRIPTION
	A string representation of the Envoy showing its path

reprlist ¶

__reprlist__()

repr ¶

__repr__()

Representation of the Envoy.

RETURNS	DESCRIPTION
	The string representation of the Envoy

getattr ¶

__getattr__(name: str) -> Union[Module, Envoy, Any]

Get an attribute from the underlying module.

If the attribute is callable, it will be wrapped in a tracer to enable intervention during execution.

PARAMETER	DESCRIPTION
`name`	The name of the attribute to get TYPE: `str`

RETURNS	DESCRIPTION
`Union[Module, Envoy, Any]`	The attribute value, possibly wrapped in a tracer

RAISES	DESCRIPTION
`AttributeError`	If the attribute doesn't exist

setattr ¶

__setattr__(key: Any, value: Any) -> None

Set an attribute on the Envoy.

If the value is a PyTorch module, it will be wrapped in an Envoy to enable intervention during execution.

Otherwise the write is mirrored to both the Envoy and the wrapped module when applicable, so that reads via __dict__ short-circuit and reads via __getattr__ (which falls through to _module) stay consistent.

PARAMETER	DESCRIPTION
`key`	The attribute name TYPE: `Any`
`value`	The attribute value TYPE: `Any`

getstate ¶

__getstate__()

setstate ¶

__setstate__(state)

Aliaser ¶

Aliaser(rename: Dict[str, Union[str, List[str]]])

PARAMETER	DESCRIPTION
`rename`	Dictionary mapping module names to alias names. Examples: {"layer1": "first_layer", "layer2": "second_layer"} Examples: {".model.layers": ".layers"} <-- Mounts .layers to the root model. Examples: {".transformer": ["model", "mdl"]} <-- Allows access of .transformer as .model or .mdl TYPE: `Dict[str, Union[str, List[str]]]`

ATTRIBUTE	DESCRIPTION
`rename`	Dictionary mapping module names to alias names. TYPE: `Dict[str, Union[str, List[str]]]`
`alias_to_name`	Dictionary mapping alias names to module names. TYPE: `Dict[str, str]`
`name_to_aliases`	Dictionary mapping module names to list of alias names. TYPE: `Dict[str, List[str]]`
`extras`	Dictionary mapping attribute paths (.transformer.h) to list of alias names. Used to show dot seperated attributes in the string representation of the Envoy. TYPE: `Dict[str, List[str]]`

rename `instance-attribute` ¶

rename = rename

alias_to_name `instance-attribute` ¶

alias_to_name = {}

name_to_aliases `instance-attribute` ¶

name_to_aliases = {}

extras `instance-attribute` ¶

extras = {}

build ¶

build(envoy: Envoy)

trace_only ¶

trace_only(fn: Callable)

envoy¶

envoy ¶

Envoy ¶

path instance-attribute ¶

interleaver instance-attribute ¶

interleaving property ¶

source property ¶

iter property ¶

device property ¶

devices property ¶

__getitem__ ¶

output ¶

inputs ¶

input ¶

__call__ ¶

trace ¶

scan ¶

edit ¶

clear_edits ¶

export_edits ¶

import_edits ¶

session ¶

all ¶

next ¶

skip ¶

to ¶

cpu ¶

cuda ¶

modules ¶

named_modules ¶

get ¶

interleave ¶

__len__ ¶

__iter__ ¶

__str__ ¶

__reprlist__ ¶

__repr__ ¶

__getattr__ ¶

__setattr__ ¶

__getstate__ ¶

__setstate__ ¶

Aliaser ¶

rename instance-attribute ¶

alias_to_name instance-attribute ¶

name_to_aliases instance-attribute ¶

extras instance-attribute ¶

build ¶

trace_only ¶

path `instance-attribute` ¶

interleaver `instance-attribute` ¶

interleaving `property` ¶

source `property` ¶

iter `property` ¶

device `property` ¶

devices `property` ¶

getitem ¶

call ¶

len ¶

iter ¶

str ¶

reprlist ¶

repr ¶

getattr ¶

setattr ¶

getstate ¶

setstate ¶

rename `instance-attribute` ¶

alias_to_name `instance-attribute` ¶

name_to_aliases `instance-attribute` ¶

extras `instance-attribute` ¶