Struct MoELayer 

pub struct MoELayer {
    pub router_weights: Tensor,
    pub num_experts: usize,
    pub num_experts_per_tok: usize,
    pub device: Device,
}

Mixture of Experts layer Routes tokens to top-k experts and combines their outputs

Fields

router_weights: Tensor

Router weights: [hidden_size, num_experts]

num_experts: usize

Number of experts

num_experts_per_tok: usize

Number of experts per token

device: Device

Device

Implementations

impl MoELayer

pub fn new( router_weights: Tensor, num_experts: usize, num_experts_per_tok: usize, device: Device, ) -> Self

Create new MoE layer

pub fn route(&self, hidden_states: &Tensor) -> Result<(Tensor, Tensor)>

Compute routing weights and indices for each token Returns (routing_weights, expert_indices) where:

• routing_weights: [batch * seq_len, num_experts_per_tok] - normalized weights
• expert_indices: [batch * seq_len, num_experts_per_tok] - expert indices

pub fn forward_with_experts<F>( &self, hidden_states: &Tensor, expert_fn: F, ) -> Result<Tensor>

where F: Fn(&Tensor, usize) -> Result<Tensor>,

Forward pass through MoE layer expert_fn: function that takes (hidden_states, expert_idx) and returns expert output

Trait Implementations

impl Debug for MoELayer

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.