pub fn gate(
x: &[f32],
w: &[f32],
bias: &[f32],
tokens: usize,
dim: usize,
n_routed: usize,
topk: usize,
gate_temp: f32,
norm_topk_prob: bool,
route_scale: f32,
) -> (Vec<f32>, Vec<usize>)Expand description
Gate.forward for score_func = "sqrtsoftplus". x: [tokens, dim] bf16 values;
w: [n_routed, dim] bf16 values; bias: [n_routed] f32. Returns
(weights[tokens, topk] f32, indices[tokens, topk]), indices in torch’s topk order
(descending by score + bias).