| 141 | return x, weight |
| 142 | |
| 143 | def ATA(self, xr=None, xl=None): |
| 144 | # xl^T * A^T * A * xr |
| 145 | assert xr is not None or xl is not None |
| 146 | weight = self._get_weight() |
| 147 | if xr is not None: |
| 148 | Axr = F.linear(input=xr, weight=weight, bias=None) # (*, cout) |
| 149 | ATAxr = F.linear(input=Axr, weight=weight.t(), bias=None) # (*, cin) |
| 150 | if xl is not None: |
| 151 | xlATAxr = torch.einsum("...i,...i->...", xl, ATAxr) # (*,) |
| 152 | return xlATAxr, weight |
| 153 | else: |
| 154 | return ATAxr, weight |
| 155 | else: |
| 156 | assert xl is not None |
| 157 | Axl = F.linear(input=xl, weight=weight, bias=None) # (*, cout) |
| 158 | ATAxl = F.linear(input=Axl, weight=weight.t(), bias=None) # (*, cin) |
| 159 | return ATAxl, weight |
| 160 | |
| 161 | def ATAATA(self, xr=None, xl=None): |
| 162 | # xl^T * A^T * A * A^T * A * xr |