r""" Args: input: - For A: `(*, in_features)` - For AT: `(*, out_features)` - For ATA: `(*, in_features)` mode (str): - ``'A'``: :math:`A(x) = y = A * x` - ``'AT'``: :math:`AT(y) = x = A^T
(
self,
input: torch.Tensor,
mode: str = "A",
input_left: torch.Tensor = None,
)
| 69 | ) |
| 70 | |
| 71 | def forward( |
| 72 | self, |
| 73 | input: torch.Tensor, |
| 74 | mode: str = "A", |
| 75 | input_left: torch.Tensor = None, |
| 76 | ): |
| 77 | r""" |
| 78 | Args: |
| 79 | input: |
| 80 | - For A: `(*, in_features)` |
| 81 | - For AT: `(*, out_features)` |
| 82 | - For ATA: `(*, in_features)` |
| 83 | mode (str): |
| 84 | - ``'A'``: :math:`A(x) = y = A * x` |
| 85 | - ``'AT'``: :math:`AT(y) = x = A^T * y` |
| 86 | - ``'ATA'``: :math:`ATA(x_r, x_l) = x_l^T * A^T * A * x_r` |
| 87 | - ``'raw'``: returns the weight |
| 88 | input_left `(*, in_features)`: |
| 89 | only used when the mode is ``'ATA'``: `(*, in_features)` |
| 90 | |
| 91 | Returns: |
| 92 | - For A: a tuple [`(*, out_features)` and weight `(out_features, in_features)`] |
| 93 | - For AT: a tuple [`(*, in_features)` and weight `(out_features, in_features)`] |
| 94 | - For ATA: a tuple [if input_left is `None`, `(*, in_features)` else `(*,)`, and weight `(out_features, in_features)`] |
| 95 | - For raw: weight `(out_features, in_features)` |
| 96 | |
| 97 | """ |
| 98 | |
| 99 | if mode == "A": |
| 100 | return self.A(input) |
| 101 | elif mode == "AT": |
| 102 | return self.AT(input) |
| 103 | elif mode == "ATA": |
| 104 | return self.ATA(xr=input, xl=input_left) |
| 105 | elif mode == "ATAATA": |
| 106 | return self.ATAATA(xr=input, xl=input_left) |
| 107 | elif mode == "raw": |
| 108 | return self._get_weight() |
| 109 | else: |
| 110 | raise NotImplementedError |
| 111 | |
| 112 | def _get_weight(self): |
| 113 | weight = self.scale * self.weight # (cout, cin) |
nothing calls this directly
no test coverage detected