(module, config)
| 48 | return (x * cos) + (rotate_every_two(x) * sin) |
| 49 | |
| 50 | def _init_weights(module, config): |
| 51 | if isinstance(module, (nn.Linear,)): |
| 52 | # Slightly different from Mesh Transformer JAX which uses truncated_normal for initialization |
| 53 | # cf https://github.com/pytorch/pytorch/pull/5617 |
| 54 | module.weight.data.normal_(mean=0.0, std=config.initializer_range) |
| 55 | if module.bias is not None: |
| 56 | module.bias.data.zero_() |
| 57 | elif isinstance(module, nn.Embedding): |
| 58 | module.weight.data.normal_(mean=0.0, std=config.initializer_range) |
| 59 | if module.padding_idx is not None: |
| 60 | module.weight.data[module.padding_idx].zero_() |
| 61 | elif isinstance(module, nn.LayerNorm): |
| 62 | module.bias.data.zero_() |
| 63 | module.weight.data.fill_(1.0) |
| 64 | |
| 65 | def _convert_head_mask_to_5d(head_mask, num_hidden_layers, dtype): |
| 66 | """-> [num_hidden_layers x batch x num_heads x seq_length x seq_length]""" |
nothing calls this directly
no outgoing calls
no test coverage detected