oa::matrix::top_k_mask

Expand I32 [T, K] expert indices into an exact FP32 [T, E] mask. Errors Returns an error unless the input is nonempty rank-two I32, num_experts is positive, all sizes fit the shader ABI, or runtime recording fails.

FnRust public API

Function Documentation

pub fn top_k_mask(
	indices: &Matrix,
	num_experts: usize
) -> Result<Matrix>

Expand I32 [T, K] expert indices into an exact FP32 [T, E] mask. Errors Returns an error unless the input is nonempty rank-two I32, num_experts is positive, all sizes fit the shader ABI, or runtime recording fails.