nnsight.util

Module for utility functions and classes used throughout the package.

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: Any, fn: Callable, cls: Type, inplace: bool = False) → Collection

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

nnsight.util.wrap(object: object, wrapper: Type, *args, **kwargs) → object

Wraps some object given some wrapper type. Updates the __class__ attribute of the object and calls the wrapper type’s __init__ method.

Parameters:

• object (object) – Object to wrap.

• wrapper (Type) – Type to wrap the object in.

Returns:

Wrapped object.

Return type:

object