oa::ml::matrix::split_heads

Permute contiguous [B*S, D] storage into [B*H, S, D/H] heads. The one-head case is a zero-copy differentiable view. Errors Returns an error unless input is nonempty F32 [B*S, D], all dimensions are nonzero, D is divisible by num_heads, or runtime recording fails.

FnRust public API

Function Documentation

pub fn split_heads( input: &Matrix, batch: usize, sequence_length: usize, num_heads: usize ) -> Result<Matrix>
Permute contiguous [B*S, D] storage into [B*H, S, D/H] heads. The one-head case is a zero-copy differentiable view. Errors Returns an error unless input is nonempty F32 [B*S, D], all dimensions are nonzero, D is divisible by num_heads, or runtime recording fails.