| 579 | } |
| 580 | |
| 581 | core::TensorValue graph_self_attention( |
| 582 | core::ModuleBuildContext & ctx, |
| 583 | const core::TensorValue & x, |
| 584 | const std::vector<core::TensorValue> & attn, |
| 585 | const std::unordered_map<std::string, Param> & params, |
| 586 | const std::string & prefix, |
| 587 | const std::string & which) { |
| 588 | const std::string base = prefix + "." + which; |
| 589 | auto v = graph_linear(ctx, x, params, base + ".in_proj", kHeads * kValueHeadDim); |
| 590 | std::vector<core::TensorValue> pieces; |
| 591 | pieces.reserve(kHeads); |
| 592 | for (int64_t h = 0; h < kHeads; ++h) { |
| 593 | auto vh = modules::SliceModule({2, h * kValueHeadDim, kValueHeadDim}).build(ctx, v); |
| 594 | vh = graph_transpose(ctx, vh, {{1, 0, 2, 3}}, 3); |
| 595 | auto out = modules::MatMulModule().build(ctx, attn[static_cast<size_t>(h)], vh); |
| 596 | pieces.push_back(graph_transpose(ctx, out, {{1, 0, 2, 3}}, 3)); |
| 597 | } |
| 598 | auto ctx_value = graph_concat_all(ctx, pieces, 2); |
| 599 | return graph_linear(ctx, ctx_value, params, base + ".out_proj", kDense); |
| 600 | } |
| 601 | |
| 602 | core::TensorValue graph_feed_forward( |
| 603 | core::ModuleBuildContext & ctx, |
no test coverage detected