gate

Function gate 

Source
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).