Coverage for transformer_lens/model_bridge/generalized_components/__init__.py: 100%
49 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Bridge components for transformer architectures."""
3from transformer_lens.model_bridge.generalized_components.alibi_joint_qkv_attention import (
4 ALiBiJointQKVAttentionBridge,
5)
6from transformer_lens.model_bridge.generalized_components.altup_block import (
7 AltUpBlockBridge,
8)
9from transformer_lens.model_bridge.generalized_components.attention import (
10 AttentionBridge,
11)
12from transformer_lens.model_bridge.generalized_components.audio_feature_extractor import (
13 AudioFeatureExtractorBridge,
14)
15from transformer_lens.model_bridge.generalized_components.block import (
16 BlockBridge,
17 DelegatedAttentionBlockBridge,
18 MLABlockBridge,
19 ParallelBlockBridge,
20 ScaledResidualBlockBridge,
21)
22from transformer_lens.model_bridge.generalized_components.bloom_attention import (
23 BloomAttentionBridge,
24)
25from transformer_lens.model_bridge.generalized_components.bloom_block import (
26 BloomBlockBridge,
27)
28from transformer_lens.model_bridge.generalized_components.bloom_mlp import (
29 BloomMLPBridge,
30)
31from transformer_lens.model_bridge.generalized_components.clip_vision_encoder import (
32 CLIPVisionEncoderBridge,
33 CLIPVisionEncoderLayerBridge,
34)
35from transformer_lens.model_bridge.generalized_components.codegen_attention import (
36 CodeGenAttentionBridge,
37)
38from transformer_lens.model_bridge.generalized_components.conv1d import Conv1DBridge
39from transformer_lens.model_bridge.generalized_components.conv_pos_embed import (
40 ConvPosEmbedBridge,
41)
42from transformer_lens.model_bridge.generalized_components.depthwise_conv1d import (
43 DepthwiseConv1DBridge,
44)
45from transformer_lens.model_bridge.generalized_components.embedding import (
46 EmbeddingBridge,
47)
48from transformer_lens.model_bridge.generalized_components.gated_delta_net import (
49 GatedDeltaNetBridge,
50)
51from transformer_lens.model_bridge.generalized_components.gated_mlp import (
52 GatedMLPBridge,
53)
54from transformer_lens.model_bridge.generalized_components.gated_rms_norm import (
55 GatedRMSNormBridge,
56)
57from transformer_lens.model_bridge.generalized_components.glm_moe_dsa_attention import (
58 GlmMoeDsaAttentionBridge,
59)
60from transformer_lens.model_bridge.generalized_components.joint_gate_up_mlp import (
61 JointGateUpMLPBridge,
62)
63from transformer_lens.model_bridge.generalized_components.joint_qkv_attention import (
64 JointQKVAttentionBridge,
65)
66from transformer_lens.model_bridge.generalized_components.joint_qkv_position_embeddings_attention import (
67 JointQKVPositionEmbeddingsAttentionBridge,
68)
69from transformer_lens.model_bridge.generalized_components.lfm2_gated_short_conv import (
70 Lfm2ShortConvBridge,
71)
72from transformer_lens.model_bridge.generalized_components.linear import LinearBridge
73from transformer_lens.model_bridge.generalized_components.mla_attention import (
74 MLAAttentionBridge,
75)
76from transformer_lens.model_bridge.generalized_components.mlp import MLPBridge
77from transformer_lens.model_bridge.generalized_components.moe import (
78 MoEBridge,
79 MoERouterBridge,
80)
81from transformer_lens.model_bridge.generalized_components.mpt_alibi_attention import (
82 MPTALiBiAttentionBridge,
83)
84from transformer_lens.model_bridge.generalized_components.normalization import (
85 NormalizationBridge,
86)
87from transformer_lens.model_bridge.generalized_components.opaque_block import (
88 OpaqueBlockBridge,
89)
90from transformer_lens.model_bridge.generalized_components.pos_embed import (
91 PosEmbedBridge,
92)
93from transformer_lens.model_bridge.generalized_components.position_embeddings_attention import (
94 PositionEmbeddingsAttentionBridge,
95)
96from transformer_lens.model_bridge.generalized_components.qwen3_5_vision_encoder import (
97 Qwen3_5VisionBlockBridge,
98 Qwen3_5VisionEncoderBridge,
99)
100from transformer_lens.model_bridge.generalized_components.rms_normalization import (
101 RMSNormalizationBridge,
102)
103from transformer_lens.model_bridge.generalized_components.rotary_embedding import (
104 RotaryEmbeddingBridge,
105)
106from transformer_lens.model_bridge.generalized_components.siglip_vision_encoder import (
107 SiglipVisionEncoderBridge,
108 SiglipVisionEncoderLayerBridge,
109)
110from transformer_lens.model_bridge.generalized_components.ssm2_mixer import (
111 SSM2MixerBridge,
112)
113from transformer_lens.model_bridge.generalized_components.ssm_block import (
114 SSMBlockBridge,
115)
116from transformer_lens.model_bridge.generalized_components.ssm_mixer import (
117 SSMMixerBridge,
118)
119from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
120 SSMMixerProtocol,
121 SSMStateHookMixin,
122 find_ssm_mixer,
123)
124from transformer_lens.model_bridge.generalized_components.symbolic import SymbolicBridge
125from transformer_lens.model_bridge.generalized_components.t5_block import T5BlockBridge
126from transformer_lens.model_bridge.generalized_components.t5gemma2_decoder_block import (
127 T5Gemma2DecoderBlockBridge,
128)
129from transformer_lens.model_bridge.generalized_components.t5gemma2_merged_attention import (
130 T5Gemma2MergedAttentionBridge,
131)
132from transformer_lens.model_bridge.generalized_components.t5gemma_decoder_block import (
133 T5GemmaDecoderBlockBridge,
134)
135from transformer_lens.model_bridge.generalized_components.unembedding import (
136 UnembeddingBridge,
137)
138from transformer_lens.model_bridge.generalized_components.vision_projection import (
139 VisionProjectionBridge,
140)
141from transformer_lens.model_bridge.generalized_components.vision_embeddings import (
142 VisionEmbeddingsBridge,
143)
144from transformer_lens.model_bridge.generalized_components.vision_classifier_head import (
145 VisionClassifierHeadBridge,
146)
148__all__ = [
149 "AttentionBridge",
150 "AudioFeatureExtractorBridge",
151 "BlockBridge",
152 "DelegatedAttentionBlockBridge",
153 "Lfm2ShortConvBridge",
154 "MLABlockBridge",
155 "ParallelBlockBridge",
156 "ScaledResidualBlockBridge",
157 "BloomBlockBridge",
158 "BloomAttentionBridge",
159 "CodeGenAttentionBridge",
160 "BloomMLPBridge",
161 "CLIPVisionEncoderBridge",
162 "CLIPVisionEncoderLayerBridge",
163 "Conv1DBridge",
164 "ConvPosEmbedBridge",
165 "DepthwiseConv1DBridge",
166 "EmbeddingBridge",
167 "ALiBiJointQKVAttentionBridge",
168 "RotaryEmbeddingBridge",
169 "PosEmbedBridge",
170 "NormalizationBridge",
171 "RMSNormalizationBridge",
172 "JointQKVAttentionBridge",
173 "JointQKVPositionEmbeddingsAttentionBridge",
174 "JointGateUpMLPBridge",
175 "LinearBridge",
176 "MLPBridge",
177 "MPTALiBiAttentionBridge",
178 "AltUpBlockBridge",
179 "GatedMLPBridge",
180 "GlmMoeDsaAttentionBridge",
181 "MLAAttentionBridge",
182 "GatedRMSNormBridge",
183 "MoEBridge",
184 "MoERouterBridge",
185 "PositionEmbeddingsAttentionBridge",
186 "Qwen3_5VisionBlockBridge",
187 "Qwen3_5VisionEncoderBridge",
188 "SymbolicBridge",
189 "UnembeddingBridge",
190 "T5BlockBridge",
191 "T5GemmaDecoderBlockBridge",
192 "T5Gemma2DecoderBlockBridge",
193 "T5Gemma2MergedAttentionBridge",
194 "SiglipVisionEncoderBridge",
195 "SiglipVisionEncoderLayerBridge",
196 "SSM2MixerBridge",
197 "SSMMixerProtocol",
198 "SSMStateHookMixin",
199 "find_ssm_mixer",
200 "OpaqueBlockBridge",
201 "SSMBlockBridge",
202 "SSMMixerBridge",
203 "VisionProjectionBridge",
204 "VisionEmbeddingsBridge",
205 "VisionClassifierHeadBridge",
206]