Custom type definitions

The library is build using the Python typing system. Based on this type system we define the following custom type definitions.


from typing import BinaryIO, Callable, Dict, List, Tuple, Union

from torch import Tensor
from import Example

ActivationName = Tuple[int, str]  # (layer, name)
ActivationNames = List[ActivationName]

ActivationDict = Dict[ActivationName, Tensor]

# LM's layer sizes: (layer, name) -> size
SizeDict = Dict[ActivationName, int]

ActivationFiles = Dict[ActivationName, BinaryIO]

# token index, corpus item -> bool
SelectionFunc = Callable[[int, Example], bool]

# [(start, stop)]
ActivationRanges = List[Tuple[int, int]]

RemoveCallback = Callable[[], None]

# Activation indexing, as done in ActivationReader
ActivationIndex = Union[int, slice, List[int], Tensor]

ActivationKey = Union[ActivationIndex, Tuple[ActivationIndex, ActivationName]]


from collections import namedtuple
from typing import Callable, Union

from import Example

DataDict = namedtuple(

DataSplit = namedtuple("DataSplit", ["activation_reader", "labels", "control_labels"])

DCConfig = namedtuple(
    ["lr", "max_epochs", "rank", "lambda1", "verbose"],

# w position, batch item -> label
ControlTask = Callable[[int, Example], Union[str, int]]