(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None)
| 59 | self.register_buffer('means', torch.zeros(num_intermediates, embed_dim, dtype=dtype)) |
| 60 | |
| 61 | def forward(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None) -> InterFeatState: |
| 62 | if rot_index is None: |
| 63 | rot_index = index |
| 64 | |
| 65 | if skip: |
| 66 | assert x.ndim == 3, f'Cannot use the `skip` parameter when the `x` tensor isn\'t 3-dimensional.' |
| 67 | prefix, x = x[:, :skip], x[:, skip:] |
| 68 | |
| 69 | rotation = self._get_rotation(rot_index) |
| 70 | y = _run_kernel(x, self.means[index], rotation) |
| 71 | |
| 72 | alpha = self.alphas[index] |
| 73 | if skip: |
| 74 | alpha = torch.cat([ |
| 75 | torch.ones(skip, dtype=alpha.dtype, device=alpha.device), |
| 76 | alpha[None].expand(y.shape[1]), |
| 77 | ]).reshape(1, -1, 1) |
| 78 | y = torch.cat([prefix, y], dim=1) |
| 79 | else: |
| 80 | if x.ndim == 3: |
| 81 | alpha = alpha.reshape(1, 1, 1).expand(1, y.shape[1], 1) |
| 82 | elif x.ndim == 4: |
| 83 | alpha = alpha.reshape(1, 1, 1, 1).expand(1, 1, *y.shape[2:]) |
| 84 | else: |
| 85 | raise ValueError(f'Unsupported input dimension: {x.ndim}') |
| 86 | |
| 87 | return InterFeatState(y, alpha) |
| 88 | |
| 89 | def _get_rotation(self, rot_index: int) -> torch.Tensor: |
| 90 | if self.rotation.ndim == 2: |
nothing calls this directly
no test coverage detected