Forward through MLP layers from state dict.
(x, sd, prefix, num_layers, activation="relu")
| 84 | |
| 85 | |
| 86 | def mlp_forward(x, sd, prefix, num_layers, activation="relu"): |
| 87 | """Forward through MLP layers from state dict.""" |
| 88 | for i in range(num_layers): |
| 89 | w = sd[f"{prefix}.layers.{i}.weight"] |
| 90 | b = sd[f"{prefix}.layers.{i}.bias"] |
| 91 | x = F.linear(x, w, b) |
| 92 | if i < num_layers - 1: |
| 93 | x = F.relu(x) |
| 94 | return x |
| 95 | |
| 96 | |
| 97 | def multihead_attention_forward(query, key, value, w_q, b_q, w_k, b_k, w_v, b_v, |