aquin.compute.feature_analysis Namespace Reference

Functions
	_get_sae_path_for_layer (str model_id, int|None layer=None)
	_get_norm_path_for_layer (str model_id, int|None layer=None)
SparseAutoencoder	_load_sae_native (str model_id, int|None layer=None)
SparseAutoencoder	load_sae (str model_id="llama-3.2-1b", int|None layer=None)
dict|None	load_norm (str model_id="llama-3.2-1b", int|None layer=None)
torch.Tensor	normalize (torch.Tensor x, str model_id="llama-3.2-1b", int|None layer=None)
torch.Tensor	_sae_feature_ablate_hook (torch.Tensor value, int pos, int feature_idx, SparseAutoencoder sae, str model_id, int sae_layer)
str	label_feature_causally (int fi, str prompt, HookedTransformer model, SparseAutoencoder sae, client, int top_k=5, str model_id="llama-3.2-1b", int|None layer=None)
str	get_causal_label (int fi, str prompt, HookedTransformer model, SparseAutoencoder sae, client, str model_id="llama-3.2-1b", int|None layer=None)
str	format_feature_ref (int feature_idx, str|None label=None)
str	prompt_for_labeling (dict|None ctx=None, dict|None args=None)
None	_attach_label_to_feature_dict (dict f, *, str prompt, HookedTransformer model, SparseAutoencoder sae, client, str model_id, int|None layer, set[int]|None seen=None)
dict	label_inspection_features (dict feat_result, *, str prompt, HookedTransformer model, client, str model_id="llama-3.2-1b", int|None layer=None)
dict	enrich_feature_tool_result (dict result, *, str prompt, HookedTransformer model, client, str model_id="llama-3.2-1b", int|None layer=None, bool label_neighbors=False)
str	resolve_feature_label (int feature_idx, *, dict ctx, dict|None args=None, int|None layer=None)
dict	_run_sae_pass (str prompt, str response, HookedTransformer model, int top_k=TOP_K_FEATURES, str model_id="llama-3.2-1b", int|None layer=None)
dict	run_feature_analysis_unlabeled (str prompt, str response, HookedTransformer model, str model_id="llama-3.2-1b", int|None layer=None)
dict	run_feature_analysis (str prompt, str response, HookedTransformer model, client, int top_k=TOP_K_FEATURES, str model_id="llama-3.2-1b", int|None layer=None)
dict	get_feature_logits (int feature_idx, HookedTransformer model, str model_id="llama-3.2-1b", int|None layer=None, int top_k=10)
dict	get_feature_neighbors (int feature_idx, str model_id="llama-3.2-1b", int|None layer=None, int top_k=8)

Variables
int	TOP_K_FEATURES = 10
	DEVICE = resolve_compute_device()
dict	_sae_cache = {}
dict	_norm_cache = {}
dict	_session_label_cache = {}
	_sae = None
	_norm = None
	_kernel_feature_acts = None
	_kernel_resid = None
list	_kernel_top_features = []

Function Documentation

_attach_label_to_feature_dict()

None _attach_label_to_feature_dict ( dict f, * , str prompt, HookedTransformer model, SparseAutoencoder sae, client, str model_id, int | None layer, set[int] | None seen = None )

protected

Definition at line 297 of file feature_analysis.py.

References format_feature_ref(), and get_causal_label().

Referenced by label_inspection_features().

_get_norm_path_for_layer()

_get_norm_path_for_layer ( str model_id, int | None layer = None )

protected

Definition at line 50 of file feature_analysis.py.

Referenced by load_norm().

_get_sae_path_for_layer()

_get_sae_path_for_layer ( str model_id, int | None layer = None )

protected

Definition at line 44 of file feature_analysis.py.

Referenced by _load_sae_native().

_load_sae_native()

SparseAutoencoder _load_sae_native ( str model_id, int | None layer = None )

protected

Definition at line 84 of file feature_analysis.py.

References _get_sae_path_for_layer().

_run_sae_pass()

dict _run_sae_pass ( str prompt, str response, HookedTransformer model, int top_k = TOP_K_FEATURES, str model_id = "llama-3.2-1b", int | None layer = None )

protected

Definition at line 420 of file feature_analysis.py.

References load_sae(), and normalize().

Referenced by run_feature_analysis(), and run_feature_analysis_unlabeled().

_sae_feature_ablate_hook()

torch.Tensor _sae_feature_ablate_hook ( torch.Tensor value, int pos, int feature_idx, SparseAutoencoder sae, str model_id, int sae_layer )

protected

Zero one SAE feature at a residual position (same logic as label_feature_causally).

Definition at line 145 of file feature_analysis.py.

References normalize().

Referenced by label_feature_causally().

enrich_feature_tool_result()

dict enrich_feature_tool_result	(	dict	result,
		*	,
		str	prompt,
		HookedTransformer	model,
			client,
		str	model_id = "llama-3.2-1b",
		int | None	layer = None,
		bool	label_neighbors = False )

Add label + feature_ref to a feature-logits or feature-neighbors payload.

Definition at line 348 of file feature_analysis.py.

References format_feature_ref(), get_causal_label(), and load_sae().

format_feature_ref()

str format_feature_ref	(	int	feature_idx,
		str | None	label = None )

Display form: index plus causal label (never index alone when label is known).

Definition at line 283 of file feature_analysis.py.

Referenced by _attach_label_to_feature_dict(), enrich_feature_tool_result(), and run_feature_analysis().

get_causal_label()

str get_causal_label	(	int	fi,
		str	prompt,
		HookedTransformer	model,
		SparseAutoencoder	sae,
			client,
		str	model_id = "llama-3.2-1b",
		int | None	layer = None )

Definition at line 269 of file feature_analysis.py.

References label_feature_causally().

Referenced by _attach_label_to_feature_dict(), enrich_feature_tool_result(), resolve_feature_label(), and run_feature_analysis().

get_feature_logits()

dict get_feature_logits	(	int	feature_idx,
		HookedTransformer	model,
		str	model_id = "llama-3.2-1b",
		int | None	layer = None,
		int	top_k = 10 )

Top vocab tokens boosted/suppressed by an SAE decoder direction (W_dec @ W_U).

Definition at line 568 of file feature_analysis.py.

References load_sae().

get_feature_neighbors()

dict get_feature_neighbors	(	int	feature_idx,
		str	model_id = "llama-3.2-1b",
		int | None	layer = None,
		int	top_k = 8 )

Cosine-nearest SAE features in decoder weight space.

Definition at line 609 of file feature_analysis.py.

References load_sae().

label_feature_causally()

str label_feature_causally	(	int	fi,
		str	prompt,
		HookedTransformer	model,
		SparseAutoencoder	sae,
			client,
		int	top_k = 5,
		str	model_id = "llama-3.2-1b",
		int | None	layer = None )

Definition at line 165 of file feature_analysis.py.

References _sae_feature_ablate_hook(), and normalize().

Referenced by get_causal_label().

label_inspection_features()

dict label_inspection_features	(	dict	feat_result,
		*	,
		str	prompt,
		HookedTransformer	model,
			client,
		str	model_id = "llama-3.2-1b",
		int | None	layer = None )

Attach causal labels to inspection feature lists (top + attribution).

Definition at line 319 of file feature_analysis.py.

References _attach_label_to_feature_dict(), and load_sae().

load_norm()

dict | None load_norm	(	str	model_id = "llama-3.2-1b",
		int | None	layer = None )

Definition at line 116 of file feature_analysis.py.

References _get_norm_path_for_layer().

Referenced by normalize().

load_sae()

SparseAutoencoder load_sae	(	str	model_id = "llama-3.2-1b",
		int | None	layer = None )

Definition at line 101 of file feature_analysis.py.

Referenced by _run_sae_pass(), enrich_feature_tool_result(), get_feature_logits(), get_feature_neighbors(), label_inspection_features(), resolve_feature_label(), and run_feature_analysis().

normalize()

torch.Tensor normalize	(	torch.Tensor	x,
		str	model_id = "llama-3.2-1b",
		int | None	layer = None )

Definition at line 135 of file feature_analysis.py.

References load_norm().

Referenced by _run_sae_pass(), _sae_feature_ablate_hook(), and label_feature_causally().

prompt_for_labeling()

str prompt_for_labeling	(	dict | None	ctx = None,
		dict | None	args = None )

Definition at line 290 of file feature_analysis.py.

Referenced by resolve_feature_label().

resolve_feature_label()

str resolve_feature_label	(	int	feature_idx,
		*	,
		dict	ctx,
		dict | None	args = None,
		int | None	layer = None )

Resolve a causal label using session context (for steer / UI tools).

Definition at line 383 of file feature_analysis.py.

References get_causal_label(), load_sae(), and prompt_for_labeling().

run_feature_analysis()

dict run_feature_analysis	(	str	prompt,
		str	response,
		HookedTransformer	model,
			client,
		int	top_k = TOP_K_FEATURES,
		str	model_id = "llama-3.2-1b",
		int | None	layer = None )

Definition at line 538 of file feature_analysis.py.

References _run_sae_pass(), format_feature_ref(), get_causal_label(), and load_sae().

run_feature_analysis_unlabeled()

dict run_feature_analysis_unlabeled	(	str	prompt,
		str	response,
		HookedTransformer	model,
		str	model_id = "llama-3.2-1b",
		int | None	layer = None )

Definition at line 534 of file feature_analysis.py.

References _run_sae_pass().

Variable Documentation

_kernel_feature_acts

aquin.compute.feature_analysis._kernel_feature_acts = None

protected

Definition at line 39 of file feature_analysis.py.

_kernel_resid

aquin.compute.feature_analysis._kernel_resid = None

protected

Definition at line 40 of file feature_analysis.py.

_kernel_top_features

list aquin.compute.feature_analysis._kernel_top_features = []

protected

Definition at line 41 of file feature_analysis.py.

_norm

aquin.compute.feature_analysis._norm = None

protected

Definition at line 37 of file feature_analysis.py.

_norm_cache

dict aquin.compute.feature_analysis._norm_cache = {}

protected

Definition at line 33 of file feature_analysis.py.

_sae

aquin.compute.feature_analysis._sae = None

protected

Definition at line 36 of file feature_analysis.py.

_sae_cache

dict aquin.compute.feature_analysis._sae_cache = {}

protected

Definition at line 32 of file feature_analysis.py.

_session_label_cache

dict aquin.compute.feature_analysis._session_label_cache = {}

protected

Definition at line 34 of file feature_analysis.py.

DEVICE

aquin.compute.feature_analysis.DEVICE = resolve_compute_device()

Definition at line 30 of file feature_analysis.py.

TOP_K_FEATURES

int aquin.compute.feature_analysis.TOP_K_FEATURES = 10

Definition at line 29 of file feature_analysis.py.