org

docs/index.rst

@@ -77,7 +77,6 @@ Install from github::
 
     train_examples
     notebooks/evaluation
-    notebooks/lazy_initialization
     notebooks/mixed_precision
     notebooks/checkpointing
     notebooks/regularization

docs/notebooks/layers_overview.ipynb

@@ -27,7 +27,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "**Handling weight initalization**\n",
+    "### Handling weight initalization\n",
     "\n",
     "Layers that contain `weight_init` or `bias_init` can accept:\n",
     "\n",
@@ -92,7 +92,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "**Handling activation functions**\n",
+    "### Handling activation functions\n",
     "\n",
     "Layers that contain `act_func` accepts:\n",
     "\n",
@@ -203,6 +203,66 @@
     "linear\n",
     "# note the dtype is f16(float16) in the repr output"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Lazy shape inference\n",
+    "\n",
+    "Lazy initialization is useful in scenarios where the dimensions of certain input features are not known in advance. For instance, when the number of neurons required for a flattened image input is uncertain, or the shape of the output from a flattened convolutional layer is not straightforward to calculate. In such cases, lazy initialization defers layers materialization until the first input.\n",
+    "\n",
+    "In `serket`, simply replace `in_features` with `None` to indicate that this layer is lazy. then materialzie the layer by functionally calling the layer. Recall that functional call - via `.at[method_name](*args, **kwargs)` _always_ returns a tuple of method output and a _new_ instance."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**Marking the layer lazy**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax\n",
+    "import serket as sk\n",
+    "\n",
+    "# 5 images from MNIST\n",
+    "x = jax.numpy.ones([5, 1, 28, 28])\n",
+    "\n",
+    "layer = sk.nn.Sequential(\n",
+    "    jax.numpy.ravel,\n",
+    "    # lazy in_features inference pass `None`\n",
+    "    sk.nn.Linear(None, 10),\n",
+    "    jax.nn.relu,\n",
+    "    sk.nn.Linear(10, 10),\n",
+    "    jax.nn.softmax,\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**Materialization by functional call**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# materialize the layer with single image\n",
+    "_, layer = layer.at[\"__call__\"](x[0])\n",
+    "# apply on batch\n",
+    "y = jax.vmap(layer)(x)\n",
+    "y.shape"
+   ]
   }
  ],
  "metadata": {