| 159 | return ATAxl, weight |
| 160 | |
| 161 | def ATAATA(self, xr=None, xl=None): |
| 162 | # xl^T * A^T * A * A^T * A * xr |
| 163 | assert xr is not None or xl is not None |
| 164 | weight = self._get_weight() |
| 165 | if xr is not None: |
| 166 | # A |
| 167 | y = F.linear(input=xr, weight=weight, bias=None) # (*, cout) |
| 168 | # AT |
| 169 | y = F.linear(input=y, weight=weight.t(), bias=None) # (*, cin) |
| 170 | # A |
| 171 | y = F.linear(input=y, weight=weight, bias=None) # (*, cout) |
| 172 | # AT |
| 173 | y = F.linear(input=y, weight=weight.t(), bias=None) # (*, cin) |
| 174 | if xl is not None: |
| 175 | y = torch.einsum("...i,...i->...", xl, y) # (*,) |
| 176 | return y, weight |
| 177 | else: |
| 178 | return y, weight |
| 179 | else: |
| 180 | assert xl is not None |
| 181 | # compute (xl^T * A^T * A * A^T * A)^T = At * A * At * A * xl |
| 182 | # A |
| 183 | y = F.linear(input=xl, weight=weight, bias=None) # (*, cout) |
| 184 | # AT |
| 185 | y = F.linear(input=y, weight=weight.t(), bias=None) # (*, cin) |
| 186 | # A |
| 187 | y = F.linear(input=y, weight=weight, bias=None) # (*, cout) |
| 188 | # AT |
| 189 | y = F.linear(input=y, weight=weight.t(), bias=None) # (*, cin) |
| 190 | return y, weight |
| 191 | |
| 192 | |
| 193 | class ModulatedSubspace(nn.Module): |