Struct MixtureOfExperts 

pub struct MixtureOfExperts<E, G>
where
    E: Estimator,
    G: GatingNetwork,
{ /* private fields */ }

Mixture of Experts ensemble

Implementations

impl<E, G> MixtureOfExperts<E, G>

where E: Estimator, G: GatingNetwork,

pub fn builder() -> MoeBuilder<E, G>

pub fn n_experts(&self) -> usize

pub fn config(&self) -> &MoeConfig

pub fn predict(&self, input: &[f32]) -> f32

pub fn predict_batch(&self, inputs: &Matrix<f32>) -> Vector<f32>

Predict for a batch of inputs.

Returns predictions and optionally the expert routing decisions.

Arguments

• inputs - Matrix of shape [n_samples, n_features]

Returns

Vector of predictions, one per input sample.

pub fn compute_load_balance_loss(&self, inputs: &Matrix<f32>) -> f32

Compute load balancing auxiliary loss.

Encourages even distribution of inputs across experts to prevent expert collapse (all inputs routed to single expert).

Switch Transformer (Fedus et al. 2021, Eq. 4-6) auxiliary loss: Loss = alpha * N * sum_i(f_i * P_i) where:

• N = number of experts
• alpha = load_balance_weight
• f_i = fraction of tokens dispatched to expert i (hard top-k)
• P_i = MEAN router softmax probability for expert i over ALL tokens (full softmax, every expert, every token — NOT top-k-only)

Arguments

• inputs - Matrix of shape [n_samples, n_features]

Returns

Load balance loss value (lower is more balanced)

pub fn expert_usage(&self, inputs: &Matrix<f32>) -> Vec<f32>

Get expert usage statistics for a batch of inputs.

Returns the fraction of inputs routed to each expert.

Arguments

• inputs - Matrix of shape [n_samples, n_features]

Returns

Vector of usage fractions, one per expert (sums to top_k).

pub fn get_routing_weights(&self, input: &[f32]) -> Vec<f32>

Get routing weights for a single input (useful for debugging/visualization).

Arguments

• input - Feature vector

Returns

Vector of gating weights for each expert.

impl<E, G> MixtureOfExperts<E, G>

where E: Estimator + Clone, G: GatingNetwork,

pub fn fit( &mut self, _x: &Matrix<f32>, _y: &Vector<f32>, ) -> Result<(), AprenderError>

Fit MoE using pre-trained experts.

This is a simple two-stage training approach:

1. Experts are assumed to be pre-trained (passed in via builder)
2. No gating training is performed (uses initial weights)

For more sophisticated training, use separate expert training followed by MoE construction.

Arguments

• _x - Training features (unused in this simple implementation)
• _y - Training labels (unused in this simple implementation)

impl<E, G> MixtureOfExperts<E, G>

where E: Estimator + Serialize + DeserializeOwned, G: GatingNetwork + Serialize + DeserializeOwned,

pub fn save<P>(&self, path: P) -> Result<(), AprenderError>

where P: AsRef<Path>,

pub fn load<P>(path: P) -> Result<MixtureOfExperts<E, G>, AprenderError>

where P: AsRef<Path>,

pub fn save_apr<P>(&self, path: P) -> Result<(), AprenderError>

where P: AsRef<Path>,

Trait Implementations

impl<E, G> Debug for MixtureOfExperts<E, G>

where E: Estimator + Debug, G: GatingNetwork + Debug,

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

Formats the value using the given formatter.