| 53 | } |
| 54 | |
| 55 | core::TensorValue reshape_heads( |
| 56 | core::ModuleBuildContext & ctx, |
| 57 | const core::TensorValue & input, |
| 58 | int64_t heads, |
| 59 | int64_t dim) { |
| 60 | auto contiguous = ensure_contiguous(ctx, input); |
| 61 | return core::reshape_tensor( |
| 62 | ctx, |
| 63 | contiguous, |
| 64 | core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim})); |
| 65 | } |
| 66 | |
| 67 | core::TensorValue repeat_kv_heads(core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t repeats) { |
| 68 | if (repeats == 1) { |
no test coverage detected