nnsight.util

Module for utility functions and classes used throughout the package.

exception nnsight.util.NNsightError(message: str, node_id: int, traceback_content: str | None = None)

NNsight Execption class for raising error during execution.

- message

error message.

Type:

str

- node_id

node id.

Type:

int

- traceback_content

traceback of the original exception being raised.

Type:

Optional[str]

class nnsight.util.Patch(parent: Any, replacement: Any, key: str)

Class representing a replacement of an attribute on a module.

obj

Object to replace.

Type:

Any

replacement

Object that replaces.

Type:

Any

parent

Module or class to replace attribute.

Type:

Any

patch() → None

Carries out the replacement of an object in a module/class.

restore() → None

Carries out the restoration of the original object on the objects module/class.

class nnsight.util.Patcher(patches: List[Patch] | None = None)

Context manager that patches from a list of Patches on __enter__ and restores the patch on __exit__.

patches

Type:

List[Patch]

add(patch: Patch) → None

Adds a Patch to the patches. Also calls .patch() on the Patch.

Parameters:

patch (Patch) – Patch to add.

class nnsight.util.WrapperModule(*args, **kwargs)

Simple torch module which passes it’s input through. Useful for hooking. If there is only one argument, returns the first element.

forward(*args, **kwargs)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

nnsight.util.apply(data: C, fn: Callable[[T], Any], cls: Type[T], inplace: bool = False) → C

Applies some function to all members of a collection of a give type (or types)

Parameters:

• data (Any) – Collection of data to apply function to.

• fn (Callable) – Function to apply.

• cls (type) – Type or Types to apply function to.

• inplace (bool) – If to apply the fn inplace. (For lists and dicts)

Returns:

Same kind of collection as data, after then fn has been applied to members of given type.

Return type:

Any

nnsight.util.fetch_attr(object: object, target: str) → Any

Retrieves an attribute from an object hierarchy given an attribute path. Levels are separated by ‘.’ e.x (transformer.h.1)

Parameters:

• object (object) – Root object to get attribute from.

• target (str) – Attribute path as ‘.’ separated string.

Returns:

Fetched attribute.

Return type:

Any