Struct MoeLayer 

pub struct MoeLayer {
    pub config: MoeConfig,
    pub router: Router,
    pub experts: Vec<Expert>,
}

Mixture of Experts layer combining a router with a set of expert networks.

Fields

config: MoeConfig

Configuration

router: Router

Router (gating network)

experts: Vec<Expert>

Expert networks

Implementations

impl MoeLayer

pub fn new(config: MoeConfig) -> Self

Create a new MoE layer from configuration.

pub fn forward(&self, input: &Array2<f32>) -> (Array2<f32>, RoutingResult)

Forward pass: route each token to top-k experts and combine outputs.

Arguments

• input - Input tensor of shape [batch_size, input_dim]

Returns

Tuple of (output tensor [batch_size, input_dim], routing result for loss computation)

pub fn balance_loss(&self, routing: &RoutingResult) -> f32

Compute the Switch Transformer-style load balancing auxiliary loss.

The balance loss encourages uniform expert utilization:

L_balance = num_experts * sum_i(f_i * P_i)

where:

• f_i = fraction of tokens dispatched to expert i
• P_i = mean routing probability for expert i

A perfectly balanced router produces L_balance = 1.0. Unbalanced routing produces L_balance > 1.0.

Arguments

• routing - The routing result from a forward pass

Returns

Scalar auxiliary loss value

Trait Implementations

impl Clone for MoeLayer

fn clone(&self) -> MoeLayer

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for MoeLayer

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

Formats the value using the given formatter.