| import torch | |
| from torch import nn | |
| class CustomLinear(nn.Linear): | |
| def __init__(self, *args, init_eye_val=0.0, is_diagonal=False, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| self.init_eye_val = init_eye_val | |
| class CustomDiagonalLinear(nn.Module): | |
| def __init__(self, d_model, bias=True, init_eye_val=0.0): | |
| super().__init__() | |
| self.init_eye_val = init_eye_val | |
| self.weight = nn.Parameter(torch.full((d_model,), init_eye_val)) | |
| self.bias = nn.Parameter(torch.zeros(d_model)) if bias else None | |
| def forward(self, input): | |
| out = input * self.weight | |
| if self.bias is not None: | |
| out += self.bias | |
| return out | |
| class Gate(nn.Module): | |
| def __init__(self, items, init_val=0.0): | |
| super().__init__() | |
| self.init_val = init_val | |
| self.gate = nn.Parameter(torch.full((items,), init_val)) | |
| def forward(self, input, dim): | |
| if input.ndim != 4: | |
| raise ValueError('input must be a 4D tensor') | |
| if not (0 <= dim <= 3): | |
| raise ValueError('dim must be 0, 1, 2, or 3') | |
| shape = [1] * 4 | |
| shape[dim] = -1 | |
| return input * self.gate.view(*shape) |