pattern_lens.consts
implements some constants and types
1"""implements some constants and types""" 2 3import re 4from typing import Literal 5 6import numpy as np 7import torch 8from jaxtyping import Float 9 10AttentionMatrix = Float[np.ndarray, "n_ctx n_ctx"] 11"type alias for attention matrix" 12 13ActivationCacheNp = dict[str, np.ndarray] 14"type alias for a cache of activations, like a transformer_lens.ActivationCache" 15 16ActivationCacheTorch = dict[str, torch.Tensor] 17"type alias for a cache of activations, like a transformer_lens.ActivationCache but without the extras. useful for when loading from an npz file" 18 19DATA_DIR: str = "attn_data" 20"default directory for attention data" 21 22ATTN_PATTERN_REGEX: re.Pattern = re.compile(r"blocks\.(\d+)\.attn\.hook_pattern") 23"regex for finding attention patterns in model state dicts" 24 25SPINNER_KWARGS: dict = dict( 26 config=dict(success="✔️ "), 27) 28"default kwargs for `muutils.spinner.Spinner`" 29 30DIVIDER_S1: str = "=" * 70 31"divider string for separating sections" 32 33DIVIDER_S2: str = "-" * 50 34"divider string for separating subsections" 35 36ReturnCache = Literal[None, "numpy", "torch"] 37"return type for a cache of activations"
AttentionMatrix =
<class 'jaxtyping.Float[ndarray, 'n_ctx n_ctx']'>
type alias for attention matrix
ActivationCacheNp =
dict[str, numpy.ndarray]
type alias for a cache of activations, like a transformer_lens.ActivationCache
ActivationCacheTorch =
dict[str, torch.Tensor]
type alias for a cache of activations, like a transformer_lens.ActivationCache but without the extras. useful for when loading from an npz file
DATA_DIR: str =
'attn_data'
default directory for attention data
ATTN_PATTERN_REGEX: re.Pattern =
re.compile('blocks\\.(\\d+)\\.attn\\.hook_pattern')
regex for finding attention patterns in model state dicts
SPINNER_KWARGS: dict =
{'config': {'success': '✔️ '}}
default kwargs for muutils.spinner.Spinner
DIVIDER_S1: str =
'======================================================================'
divider string for separating sections
DIVIDER_S2: str =
'--------------------------------------------------'
divider string for separating subsections
ReturnCache =
typing.Literal[None, 'numpy', 'torch']
return type for a cache of activations