oa::ml::matrix::moe_route_weights

Normalize probabilities over the selected experts for each token. probabilities is F32 [T, E]; expert_indices is I32 [T, K] with 1 <= K <= E. The routing decision is detached from reverse mode, while gradients propagate to the selected probability entries. Errors Returns an error for incompatible shapes, dtypes, engine ownership, shader ABI overflow, allocation failure, or runtime recording failure.

FnRust public API

Function Documentation

pub fn moe_route_weights( probabilities: &Matrix, expert_indices: &Matrix ) -> Result<Matrix>
Normalize probabilities over the selected experts for each token. probabilities is F32 [T, E]; expert_indices is I32 [T, K] with 1 <= K <= E. The routing decision is detached from reverse mode, while gradients propagate to the selected probability entries. Errors Returns an error for incompatible shapes, dtypes, engine ownership, shader ABI overflow, allocation failure, or runtime recording failure.