transformer_lens.model_bridge.supported_architectures.lfm2_moe module

LiquidAI LFM2 MoE architecture adapter.

class transformer_lens.model_bridge.supported_architectures.lfm2_moe.Lfm2MoeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for LiquidAI Lfm2 MoE models.

__init__(cfg: Any) → None

Initialize the Lfm2 MoE architecture adapter.

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

Set up model-specific references for component testing.

class transformer_lens.model_bridge.supported_architectures.lfm2_moe.Lfm2MoeGateBridge(*args: Any, logits_index: int = 0, **kwargs: Any)

Bases: MoERouterBridge

get_random_inputs(batch_size: int = 2, seq_len: int = 8, device: device | None = None, dtype: dtype | None = None) → Dict[str, Any]

Random inputs for router component testing.

The router runs on the reshaped [N, d_model] hidden states and takes a second expert_bias arg (use_expert_bias=True); its top-k gather is hardcoded to dim=1, so the input must be 2D or the gather indexes the sequence axis out of bounds.

Parameters:

  • batch_size – Batch size for generated inputs

  • seq_len – Sequence length for generated inputs

  • device – Device to place tensors on

  • dtype – Dtype for generated tensors (defaults to float32)

Returns:

Dictionary of input tensors matching the component’s expected input signature