RealmOA Library

oa::ml::nn::Moe

Sparse top-k mixture-of-experts feed-forward module. The implementation preserves OA’s expert-major route plan and stacked parameter layout. Routing decisions remain detached while gradients flow through selected gate magnitudes, expert projections, RMSNorm, and the residual input.

StRust public API

Declaration

pub struct Moe { /* private fields */ }
Public Rust declaration.

Public Methods

pub fn forward(&self, input: &Matrix) -> Result<Matrix>
pub const fn model_width(&self) -> usize
pub const fn hidden_width(&self) -> usize
pub const fn num_experts(&self) -> usize
pub const fn experts_per_token(&self) -> usize
pub fn set_sparse_execution(&self, sparse: bool)
pub fn sparse_execution(&self) -> bool
pub fn num_shared_experts(&self) -> usize
pub const fn epsilon(&self) -> f32
pub fn set_balance_rate(&self, gamma: f32)
pub fn balance_rate(&self) -> f32
pub fn routing_bias(&self) -> Matrix
pub fn set_aux_loss_alpha(&self, alpha: f32)
pub fn aux_loss_alpha(&self) -> f32
pub fn set_router_z_loss_beta(&self, beta: f32)
pub fn router_z_loss_beta(&self) -> f32
pub fn aux_loss(&self) -> Matrix
pub fn update_routing_bias(&self) -> Result<()>
pub fn last_selection_mask(&self) -> Option<Matrix>
pub fn last_gate_probabilities(&self) -> Option<Matrix>
pub fn route_stats(&self) -> Result<MoeRouteStats>
pub fn expert_gate_up_weight(&self) -> Parameter
pub fn expert_gate_up_bias(&self) -> Parameter
pub fn expert_down_weight(&self) -> Parameter
pub fn expert_down_bias(&self) -> Parameter
fn forward(&self, input: &Matrix) -> Result<Matrix>
fn registry(&self) -> &ModuleRegistry
fn parameters(&self) -> Vec<Parameter>
fn named_parameters(&self) -> Vec<NamedParameter>
fn buffers(&self) -> Vec<Matrix>
fn named_buffers(&self) -> Vec<NamedBuffer>
fn all_parameters(&self) -> Result<Vec<Parameter>>
fn all_named_parameters(&self) -> Result<Vec<NamedParameter>>
fn all_named_buffers(&self) -> Result<Vec<NamedBuffer>>
fn all_named_state_u32(&self) -> Result<Vec<NamedStateU32>>
fn num_parameters(&self) -> Result<usize>
fn train(&self, training: bool)
fn eval(&self)
fn is_training(&self) -> bool
fn scoped_eval(&self) -> ScopedEval<'_>

Associated Functions

pub fn with_seed( engine: &Engine, model_width: usize, hidden_width: usize, num_experts: usize, experts_per_token: usize, epsilon: f32, seed: u64, ) -> Result<Self>
pub fn with_seed_and_shared_experts( engine: &Engine, model_width: usize, hidden_width: usize, num_experts: usize, experts_per_token: usize, epsilon: f32, num_shared_experts: usize, seed: u64, ) -> Result<Self>
pub fn from_matrices( norm_weight: Matrix, router_weight: Matrix, router_bias: Matrix, gate_up_weight: Matrix, gate_up_bias: Matrix, down_weight: Matrix, down_bias: Matrix, experts_per_token: usize, epsilon: f32, ) -> Result<Self>

Public Method Documentation

pub fn forward( &self, input: &Matrix ) -> Result<Matrix>
No public source comment is attached to this declaration.
pub const fn model_width( &self ) -> usize
No public source comment is attached to this declaration.
pub const fn hidden_width( &self ) -> usize
No public source comment is attached to this declaration.
pub const fn num_experts( &self ) -> usize
No public source comment is attached to this declaration.
pub const fn experts_per_token( &self ) -> usize
No public source comment is attached to this declaration.
pub fn set_sparse_execution( &self, sparse: bool )
No public source comment is attached to this declaration.
pub fn sparse_execution( &self ) -> bool
No public source comment is attached to this declaration.
pub fn num_shared_experts( &self ) -> usize
No public source comment is attached to this declaration.
pub const fn epsilon( &self ) -> f32
No public source comment is attached to this declaration.
pub fn set_balance_rate( &self, gamma: f32 )
No public source comment is attached to this declaration.
pub fn balance_rate( &self ) -> f32
No public source comment is attached to this declaration.
pub fn routing_bias( &self ) -> Matrix
No public source comment is attached to this declaration.
pub fn set_aux_loss_alpha( &self, alpha: f32 )
No public source comment is attached to this declaration.
pub fn aux_loss_alpha( &self ) -> f32
No public source comment is attached to this declaration.
pub fn set_router_z_loss_beta( &self, beta: f32 )
No public source comment is attached to this declaration.
pub fn router_z_loss_beta( &self ) -> f32
No public source comment is attached to this declaration.
pub fn aux_loss( &self ) -> Matrix
No public source comment is attached to this declaration.
pub fn update_routing_bias(&self) -> Result<()>
No public source comment is attached to this declaration.
pub fn last_selection_mask( &self ) -> Option<Matrix>
No public source comment is attached to this declaration.
pub fn last_gate_probabilities( &self ) -> Option<Matrix>
No public source comment is attached to this declaration.
pub fn route_stats( &self ) -> Result<MoeRouteStats>
No public source comment is attached to this declaration.
pub fn expert_gate_up_weight( &self ) -> Parameter
No public source comment is attached to this declaration.
pub fn expert_gate_up_bias( &self ) -> Parameter
No public source comment is attached to this declaration.
pub fn expert_down_weight( &self ) -> Parameter
No public source comment is attached to this declaration.
pub fn expert_down_bias( &self ) -> Parameter
No public source comment is attached to this declaration.
fn forward( &self, input: &Matrix ) -> Result<Matrix>
No public source comment is attached to this declaration.
fn registry( &self ) -> &ModuleRegistry
No public source comment is attached to this declaration.
fn parameters( &self ) -> Vec<Parameter>
No public source comment is attached to this declaration.
fn named_parameters( &self ) -> Vec<NamedParameter>
No public source comment is attached to this declaration.
fn buffers( &self ) -> Vec<Matrix>
No public source comment is attached to this declaration.
fn named_buffers( &self ) -> Vec<NamedBuffer>
No public source comment is attached to this declaration.
fn all_parameters( &self ) -> Result<Vec<Parameter>>
No public source comment is attached to this declaration.
fn all_named_parameters( &self ) -> Result<Vec<NamedParameter>>
No public source comment is attached to this declaration.
fn all_named_buffers( &self ) -> Result<Vec<NamedBuffer>>
No public source comment is attached to this declaration.
fn all_named_state_u32( &self ) -> Result<Vec<NamedStateU32>>
No public source comment is attached to this declaration.
fn num_parameters( &self ) -> Result<usize>
No public source comment is attached to this declaration.
fn train( &self, training: bool )
No public source comment is attached to this declaration.
fn eval( &self )
No public source comment is attached to this declaration.
fn is_training( &self ) -> bool
No public source comment is attached to this declaration.
fn scoped_eval( &self ) -> ScopedEval<'_>
No public source comment is attached to this declaration.
pub fn with_seed( engine: &Engine, model_width: usize, hidden_width: usize, num_experts: usize, experts_per_token: usize, epsilon: f32, seed: u64 ) -> Result<Self>
No public source comment is attached to this declaration.
pub fn with_seed_and_shared_experts( engine: &Engine, model_width: usize, hidden_width: usize, num_experts: usize, experts_per_token: usize, epsilon: f32, num_shared_experts: usize, seed: u64 ) -> Result<Self>
No public source comment is attached to this declaration.
pub fn from_matrices( norm_weight: Matrix, router_weight: Matrix, router_bias: Matrix, gate_up_weight: Matrix, gate_up_bias: Matrix, down_weight: Matrix, down_bias: Matrix, experts_per_token: usize, epsilon: f32 ) -> Result<Self>
No public source comment is attached to this declaration.