(self, spec, decoder)
| 348 | self.set_layer_norm(layer_spec.ffn.layer_norm, layer.final_layer_norm) |
| 349 | |
| 350 | def set_decoder(self, spec, decoder): |
| 351 | self.set_common_layers(spec, decoder) |
| 352 | |
| 353 | for layer_spec, layer in zip(spec.layer, decoder.layers): |
| 354 | self.set_attention( |
| 355 | layer_spec.self_attention, |
| 356 | layer.self_attn, |
| 357 | self_attention=True, |
| 358 | ) |
| 359 | self.set_layer_norm( |
| 360 | layer_spec.self_attention.layer_norm, |
| 361 | layer.self_attn_layer_norm, |
| 362 | ) |
| 363 | |
| 364 | if hasattr(layer, "encoder_attn"): |
| 365 | self.set_attention( |
| 366 | layer_spec.attention, |
| 367 | layer.encoder_attn, |
| 368 | self_attention=False, |
| 369 | ) |
| 370 | self.set_layer_norm( |
| 371 | layer_spec.attention.layer_norm, |
| 372 | layer.encoder_attn_layer_norm, |
| 373 | ) |
| 374 | |
| 375 | self.set_linear(layer_spec.ffn.linear_0, layer.fc1) |
| 376 | self.set_linear(layer_spec.ffn.linear_1, layer.fc2) |
| 377 | self.set_layer_norm(layer_spec.ffn.layer_norm, layer.final_layer_norm) |
| 378 | |
| 379 | def set_attention(self, spec, attention, self_attention=False): |
| 380 | split_layers = [common_spec.LinearSpec() for _ in range(3)] |
no test coverage detected