util¶

util ¶

Traversal helpers for the arbitrary values a model hands us.

Interventions deal in whatever a forward happens to return: a tensor, a tuple of them, a dataclass output holding a list of tuples. Every layer of nnsight ends up needing the same two things from such a value — visit each tensor in it, or rebuild it with each tensor replaced — without knowing its shape in advance. That is apply. Beside it live to_import_path and from_import_path, which name a class by where it can be imported from — the form a type takes when it has to cross a serialization boundary.

The rule apply follows is that container nesting is free and object nesting is not: lists, tuples, dicts, sets and slices are traversed without limit and rebuilt by their own type (so a namedtuple stays a namedtuple), while stepping into a generic object spends one of n levels, because an arbitrary object graph has no bottom.

apply ¶

apply(data: Any, fn: Callable[[Any], Any], cls: type, n: int = 0) -> Any

Recursively apply fn to every instance of cls found in data.

Containers — list, tuple (incl. namedtuples), dict, set, frozenset and slice — are traversed without limit and rebuilt with the transformed values (the originals are left untouched). Generic objects (anything exposing a __dict__) are also descended into, with their attributes transformed and written back in place, but only up to n levels of object nesting so we never recurse without bound through an arbitrary object graph.

Container nesting is free — it does not count against n. Only stepping into a generic object spends a level.

PARAMETER	DESCRIPTION
`data`	The structure to traverse. TYPE: `Any`
`fn`	Called on each `cls` instance; its return value replaces it. TYPE: `Callable[[Any], Any]`
`cls`	The leaf type to transform. TYPE: `type`
`n`	Maximum depth of generic-object descent. `n=0` (default) disables object descent entirely — only containers are traversed. `n=1` descends into the first object encountered along a path but not into objects nested within it, and so on. TYPE: `int` DEFAULT: `0`

RETURNS	DESCRIPTION
`Any`	The transformed structure. Containers are new objects; generic objects
`Any`	(when descended into) are mutated in place and returned as-is.

to_import_path ¶

to_import_path(obj: type) -> str

The dotted module.QualName path that from_import_path resolves.

from_import_path ¶

from_import_path(path: str) -> Any

Resolve a dotted module.QualName path (as produced by to_import_path).