transformer_lens.model_bridge.supported_architectures.jetmoe module

JetMoE architecture adapter.

MIT-IBM’s JetMoE (JetMoeForCausalLM, native in transformers): the only open at-scale Mixture-of-Attention-heads model — attention Q and output projections are per-expert parallel 3D tensors behind a top-k router (experts: JetMoeMoA), with a shared fused KV projection, alongside a conventional parallel-experts MoE MLP. Both routers are hookable; the mixers delegate to HF (per-expert 3D projections have no uniform reconstruction, so no fold target exists either).

class transformer_lens.model_bridge.supported_architectures.jetmoe.JetMoeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for JetMoeForCausalLM models.

__init__(cfg: Any) → None

Initialize the JetMoE architecture adapter.

component_mapping: ComponentMapping | None

component_test_skip_suffixes: tuple = ('mlp.gate', 'attn.experts.router')

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Delegated attention computes rotary inside HF; nothing to wire.

supports_fold_ln = False

uses_split_attention: bool

weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None