Coverage for transformer_lens/tools/analysis/__init__.py: 100%

1"""Analysis tools for TransformerLens. 

2 

3This subpackage collects high-level, single-call interpretability analyses that 

4sit on top of the hook/cache system. Model support is documented per tool; 

5new analyses may target the ``TransformerBridge`` API exclusively. 

6 

7Entry points and the classes they accept are exported here; result dataclasses stay in the 

8module that returns them. 

9 

10Tools: 

11 - attribution_patching: Attribution patching (gradient-linearized activation 

12 patching) over residual-stream nodes and edges: typed computational graph, 

13 a names-filtered manual-backward gradient cache, and signed node and edge 

14 scores (EAP), plus forward-only ablate-outside faithfulness with corrupt 

15 or held-out mean replacements and circuit-relative class diagnostics. 

16 Integrated gradients (EAP-IG) are not implemented. 

17 - direct_logit_attribution: Direct Logit Attribution (DLA) over components, 

18 layers, or attention heads. 

19 - get_act_patch_direct_path, get_act_patch_direct_path_all_sources: 

20 re-exported from ``transformer_lens.patching``; exact direct path 

21 patching for head-to-head circuit analysis. 

22 - interchange: Interchange interventions on linear subspaces: set a base's 

23 activation inside a subspace to its value on a source input, with 

24 difference-of-means, probe, random and neuron bases, a norm-matched 

25 placebo, batched gradient-free evaluation, distributed alignment search 

26 (DAS) to learn the subspace, and DAS's trained controls (arbitrary labels 

27 and a randomly initialised model). 

28 - jacobian_lens: The Jacobian lens (J-lens) — per-layer causal transport to 

29 the output vocabulary basis, with loading of published lens artifacts, 

30 native fitting, readouts, interventions, J-space sparse decomposition, and 

31 anchored coordinate patching (offline and dynamic/hooked). 

32 - logit_lens: One-call logit lens over the accumulated residual stream and 

33 the ``logit_readout`` kernel that reads any ``d_model`` vector through the 

34 real final norm, unembedding and adapter output transform, chunked over 

35 the vocabulary with target ranks, top-k and vocabulary subsets. 

36 - projection_kernel: Basis-invariant subspace overlap and TransformerBridge 

37 attention-head OQ/OK/OV affinity. 

38 - sparse_probing: Leakage-safe k-sparse binary probes over supplied 

39 activation tensors, with train-only selection and raw null controls. 

40""" 

41 

42import warnings 

43from typing import TYPE_CHECKING, Any 

44 

45from transformer_lens.patching import ( 

46 get_act_patch_direct_path, 

47 get_act_patch_direct_path_all_sources, 

48) 

49from transformer_lens.tools.analysis.attribution_patching import ( 

50 EdgeAttributionConfig, 

51 FaithfulnessConfig, 

52 Node, 

53 attribution_patch, 

54 faithfulness, 

55) 

56from transformer_lens.tools.analysis.direct_logit_attribution import ( 

57 direct_logit_attribution, 

58) 

59from transformer_lens.tools.analysis.interchange import ( 

60 LearnedSubspace, 

61 NormMatchedPlacebo, 

62 SubspaceInterchange, 

63 arbitrary_targets, 

64 collect_activations, 

65 fit_das, 

66 intervened_logits, 

67 mean_difference_basis, 

68 neuron_basis, 

69 probe_basis, 

70 random_basis, 

71 random_init_bridge, 

72 random_neuron_basis, 

73 reinitialized, 

74 run_with_intervention, 

75) 

76from transformer_lens.tools.analysis.jacobian_lens import JacobianLens 

77from transformer_lens.tools.analysis.jacobian_lens_coordinate_patch import ( 

78 solve_coordinate_patch, 

79) 

80from transformer_lens.tools.analysis.jacobian_lens_decomposition import ( 

81 JSpaceDecomposition, 

82 get_sparse_decomposition, 

83) 

84from transformer_lens.tools.analysis.logit_lens import logit_lens, logit_readout 

85from transformer_lens.tools.analysis.projection_kernel import ( 

86 SubspaceBasis, 

87 attention_head_subspace_affinity, 

88 orthonormal_subspace, 

89 projection_kernel, 

90 random_projection_kernel_moments, 

91) 

92from transformer_lens.tools.analysis.sparse_probing import ( 

93 SparseProbeConvergenceError, 

94 fit_sparse_probe, 

95 sweep_sparse_probe, 

96) 

97 

98__all__ = [ 

99 "EdgeAttributionConfig", 

100 "FaithfulnessConfig", 

101 "JSpaceDecomposition", 

102 "JacobianLens", 

103 "LearnedSubspace", 

104 "Node", 

105 "NormMatchedPlacebo", 

106 "SparseProbeConvergenceError", 

107 "SubspaceBasis", 

108 "SubspaceInterchange", 

109 "arbitrary_targets", 

110 "attention_head_subspace_affinity", 

111 "attribution_patch", 

112 "collect_activations", 

113 "direct_logit_attribution", 

114 "faithfulness", 

115 "fit_das", 

116 "fit_sparse_probe", 

117 "get_act_patch_direct_path", 

118 "get_act_patch_direct_path_all_sources", 

119 "get_sparse_decomposition", 

120 "intervened_logits", 

121 "logit_lens", 

122 "logit_readout", 

123 "mean_difference_basis", 

124 "neuron_basis", 

125 "orthonormal_subspace", 

126 "probe_basis", 

127 "projection_kernel", 

128 "random_basis", 

129 "random_init_bridge", 

130 "random_neuron_basis", 

131 "random_projection_kernel_moments", 

132 "reinitialized", 

133 "run_with_intervention", 

134 "solve_coordinate_patch", 

135 "sweep_sparse_probe", 

136] 

137 

138# Moved to transformer_lens.utilities with unchanged results. 

139_MOVED_TO_UTILITIES = ("success_rate_ci",) 

140 

141if TYPE_CHECKING: 

142 from transformer_lens.utilities import success_rate_ci 

143else: 

144 # Hidden from type checkers, which would otherwise type every unknown name here as Any. 

145 def __getattr__(name: str) -> Any: 

146 if name not in _MOVED_TO_UTILITIES: 

147 raise AttributeError(f"module {__name__!r} has no attribute {name!r}") 

148 warnings.warn( 

149 f"{__name__}.{name} is deprecated; import it from transformer_lens.utilities.", 

150 DeprecationWarning, 

151 stacklevel=2, 

152 ) 

153 from transformer_lens import utilities 

154 

155 return getattr(utilities, name)