include project in initalize

Author: jxbz

examples/hello-mnist.ipynb

@@ -179,7 +179,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8a432072d1bb4a39b5438417ac15d5f5",
+       "model_id": "d26bcf36d33e4b108018d97c00b975fc",
        "version_major": 2,
        "version_minor": 0
       },
@@ -202,8 +202,8 @@
     "\n",
     "error = SquareError()\n",
     "\n",
-    "w = mlp.initialize(jax.random.PRNGKey(0))\n",
-    "w = mlp.project(w)\n",
+    "key = jax.random.PRNGKey(0)\n",
+    "w = mlp.initialize(key)\n",
     "    \n",
     "progress_bar = tqdm(range(steps), desc=f\"Loss: {0:.4f}\")\n",
     "for step in progress_bar:\n",
@@ -248,7 +248,7 @@
      "output_type": "stream",
      "text": [
       "Accuracy on shown samples: 5/5\n",
-      "Overall test accuracy: 97.58%\n"
+      "Overall test accuracy: 97.52%\n"
      ]
     }
    ],

examples/hello-world.ipynb

@@ -117,15 +117,15 @@
      "output_type": "stream",
      "text": [
       "Step   0 \t Loss 0.976274\n",
-      "Step 100 \t Loss 0.001985\n",
-      "Step 200 \t Loss 0.001541\n",
-      "Step 300 \t Loss 0.001189\n",
-      "Step 400 \t Loss 0.000884\n",
-      "Step 500 \t Loss 0.000625\n",
-      "Step 600 \t Loss 0.000413\n",
-      "Step 700 \t Loss 0.000251\n",
-      "Step 800 \t Loss 0.000130\n",
-      "Step 900 \t Loss 0.000049\n"
+      "Step 100 \t Loss 0.001989\n",
+      "Step 200 \t Loss 0.001537\n",
+      "Step 300 \t Loss 0.001194\n",
+      "Step 400 \t Loss 0.000885\n",
+      "Step 500 \t Loss 0.000627\n",
+      "Step 600 \t Loss 0.000420\n",
+      "Step 700 \t Loss 0.000255\n",
+      "Step 800 \t Loss 0.000134\n",
+      "Step 900 \t Loss 0.000053\n"
      ]
     }
    ],
@@ -135,7 +135,6 @@
     "\n",
     "key = jax.random.PRNGKey(0)\n",
     "w = mlp.initialize(key)\n",
-    "w = mlp.project(w)\n",
     "\n",
     "for step in range(steps):\n",
     "    # compute outputs and activations\n",

modula/atom.py

@@ -32,7 +32,7 @@ def backward(self, w, acts, grad_output):
 
     def initialize(self, key):
         weight = jax.random.normal(key, shape=(self.fanout, self.fanin))
-        return [weight]
+        return self.project([weight])
 
     def project(self, w):
         weight = w[0]