nice_hooks.activationcache

Module Contents

Classes

ActivationCache

A wrapper around a dictionary of cached activations from a model run.

Attributes

ActivationCacheLike

Any object that a ActivationCache can be constructed from.

class nice_hooks.activationcache.ActivationCache

Bases: dict[str, torch.Tensor]

A wrapper around a dictionary of cached activations from a model run. Supports similar operations to t.Tensor, which are usually applied elementwise. See ActivationCache

to(device)

Moves each tensor of the cache to a device.

Parameters:

device (Union[str, torch.device]) –

Return type:

ActivationCache

index(key)

Applies t[key] for every tensor in the activation cache. As the tensors may have different shapes, this operation only really makes sense for manipulating the batch dimension.

Parameters:

key (Union[None, _int, slice, torch.Tensor, List, Tuple]) –

Return type:

ActivationCache

nice_hooks.activationcache.ActivationCacheLike

Any object that a ActivationCache can be constructed from.