Mojo function

gather_reduce

gather_reduce[type: DType, gather_axis: Int, reduce_axis: Int, simd_width: Int, reduce_fn: fn[DType, Int](SIMD[$0, $1], SIMD[$0, $1]) -> SIMD[$0, $1], output_rank: Int, output_shape: DimList, input_rank: Int, input_shape: DimList, indices_rank: Int, indices_shape: DimList](output: NDBuffer[type, output_rank, origin, output_shape], input: NDBuffer[type, input_rank, origin, input_shape], indices: NDBuffer[int32, indices_rank, origin, indices_shape], reduce_init: SIMD[type, 1])

Computes output[i, j, k] = input[indices[i, j], k] and simultaneously reduces the output accross axis 1 to produce output[i, k].

The motivating use-case for this is multi-hot embeddings in recommender models. This provides similar functionality to Torch's EmbeddingBag layer. In that context, i is the batch dimension, j is the multi-hot dimension, and k is the embedding dimension.