Repair parametrization notebook examples

Author: lezcano

examples/parametrisations.ipynb

@@ -67,7 +67,7 @@
    "outputs": [],
    "source": [
     "class SymmetricParametrization(nn.Module):\n",
-    "    def forward(X):\n",
+    "    def forward(self, X):\n",
     "        A = X.triu()\n",
     "        return A + A.T "
    ]
@@ -113,9 +113,10 @@
    "source": [
     "class PruningParametrization(nn.Module):\n",
     "    def __init__(self, X, p_drop=0.2):\n",
+    "        super().__init__()\n",
     "        # sample zeros with probability p_drop\n",
     "        mask = torch.full_like(X, 1.0 - p_drop)\n",
-    "        self.mask = torch.bernoulli(mask)\n",
+    "        self.register_buffer(\"mask\", torch.bernoulli(mask))\n",
     "\n",
     "    def forward(self, X):\n",
     "        return X * self.mask"
@@ -322,6 +323,8 @@
     "# In other words, right_inverse is a right inverse of forward.\n",
     "\n",
     "model = nn.Linear(5, 5)\n",
+    "with torch.no_grad():\n",
+    "    model.weight.copy_(model.weight - model.weight.T)\n",
     "P.register_parametrization(model, \"weight\", Skew())\n",
     "# Just computes `model.weight` once\n",
     "with P.cached():\n",
@@ -347,11 +350,11 @@
     "    def forward(self, A):\n",
     "        # Cayley map: (I + A)(I - A)^{-1}\n",
     "        # This is orthogonal whenever A is skew-symmetric\n",
-    "        Id = torch.eye(A.size(0))\n",
+    "        Id = torch.eye(A.size(0), dtype=A.dtype, device=A.device)\n",
     "        return self.B @ torch.linalg.solve(Id - A, Id + A)\n",
     "\n",
     "    def is_orthogonal(self, X):\n",
-    "        Id = torch.eye(X.size(0))\n",
+    "        Id = torch.eye(X.size(0), dtype=X.dtype, device=X.device)\n",
     "        return torch.allclose(X.T @ X, Id, atol=1e-6)\n",
     "\n",
     "    def right_inverse(self, X):\n",
@@ -362,9 +365,14 @@
     "        return torch.zeros_like(X)\n",
     "\n",
     "\n",
-    "model = nn.Linear(5,5)\n",
+    "model = nn.Linear(4, 4)\n",
+    "X = torch.zeros_like(model.weight)\n",
+    "X.diagonal(offset=1)[::2] = 1.0\n",
+    "X = X - X.T\n",
+    "with torch.no_grad():\n",
+    "    model.weight.copy_(X)\n",
     "P.register_parametrization(model, \"weight\", Skew())\n",
-    "P.register_parametrization(model, \"weight\", Orthogonal(5))\n",
+    "P.register_parametrization(model, \"weight\", Orthogonal(4))\n",
     "\n",
     "# Sample an orthogonal matrix and initialise the layer\n",
     "X = torch.empty_like(model.weight)\n",