| 9 | |
| 10 | class TestSubspace(unittest.TestCase): |
| 11 | def _test(self, in_features, out_features, normalize, orthogonalize, init_norm): |
| 12 | net = cdsnn.Subspace( |
| 13 | in_features=in_features, |
| 14 | out_features=out_features, |
| 15 | normalize=normalize, |
| 16 | orthogonalize=orthogonalize, |
| 17 | init_norm=init_norm, |
| 18 | ) |
| 19 | |
| 20 | A = net(None, 'raw').clone() |
| 21 | xr = torch.randn(in_features) |
| 22 | xr2 = torch.randn(in_features) |
| 23 | xl = torch.randn(out_features) |
| 24 | |
| 25 | # A * xr |
| 26 | gt = A @ xr |
| 27 | y, A2 = net(xr, 'A') |
| 28 | assert torch.allclose(gt, y) |
| 29 | assert torch.allclose(A, A2) |
| 30 | |
| 31 | # A^T * xl |
| 32 | gt = A.t() @ xl |
| 33 | y, A2 = net(xl, 'AT') |
| 34 | assert torch.allclose(gt, y) |
| 35 | assert torch.allclose(A, A2) |
| 36 | |
| 37 | # x^l * A^T * A * xr |
| 38 | gt = (xr2.t() @ A.t()) @ (A @ xr) |
| 39 | y, A2 = net(xr, 'ATA', xr2) |
| 40 | assert torch.allclose(gt, y) |
| 41 | assert torch.allclose(A, A2) |
| 42 | |
| 43 | # xr^T * A^T * A * A^T * A * xr |
| 44 | gt = ((xr2.t() @ A.t()) @ A) @ (A.t() @ A @ xr) |
| 45 | y, A2 = net(xr, 'ATAATA', xr2) |
| 46 | assert torch.allclose(gt, y) |
| 47 | assert torch.allclose(A, A2) |
| 48 | |
| 49 | def test_1(self): |
| 50 | self._test( |