Function weighted_sum 

pub fn weighted_sum(x: &Array, w: &Array) -> Array

The experts’ outputs x ([..., k, D]) weighted by w ([..., k]) and summed over the k experts: (x * w[..., None]).sum(-2), products and sum in x’s dtype as mlx-lm’s MoE blocks compute them. On the GPU it is one kernel equal to that composition bit for bit when both have one floating-point dtype and k < 32 (otherwise the composition).

use nuvec_core::{array, fast};

let x = array![[[1.0f32, 2.0], [3.0, 4.0]]];
let w = array![[0.5f32, 2.0]];
assert_eq!(fast::weighted_sum(&x, &w).to_vec::<f32>().unwrap(), [6.5, 9.0]);