| 221 | # concatenate ``Skew`` and a parametrization that implements the Cayley map to get a layer with |
| 222 | # orthogonal weights |
| 223 | class CayleyMap(nn.Module): |
| 224 | def __init__(self, n): |
| 225 | super().__init__() |
| 226 | self.register_buffer("Id", torch.eye(n)) |
| 227 | |
| 228 | def forward(self, X): |
| 229 | # (I + X)(I - X)^{-1} |
| 230 | return torch.linalg.solve(self.Id - X, self.Id + X) |
| 231 | |
| 232 | layer = nn.Linear(3, 3) |
| 233 | parametrize.register_parametrization(layer, "weight", Skew()) |