oa::ml::matrix::merge_heads

Invert split_heads into contiguous [B*S, H*(D/H)] storage. The one-head case is a zero-copy differentiable view. Errors Returns an error unless input is nonempty F32 [B*H, S, D/H] with the supplied nonzero dimensions, or runtime recording fails.

FnRust public API

Function Documentation

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