docs button

docs/conf.py

@@ -136,6 +136,20 @@
 nbsphinx_codecell_lexer = "ipython"
 nbsphinx_kernel_name = "python"
 nbsphinx_timeout = 180
+# source: https://github.com/NTT123/pax/blob/main/docs/conf.py
+nbsphinx_prolog = r"""
+{% set docname = 'docs/' + env.doc2path(env.docname, base=None) %}
+
+.. only:: html
+
+    .. role:: raw-html(raw)
+        :format: html
+
+    .. nbinfo::
+
+        Open in colab:
+        :raw-html:`<a href="https://colab.research.google.com/github/asem000/serket/blob/main/{{ docname }}"><img alt="Open In Colab" src="https://colab.research.google.com/assets/colab-badge.svg" style="vertical-align:text-bottom"></a>`
+"""
 
 # Tell sphinx-autodoc-typehints to generate stub parameter annotations including
 # types, even if the parameters aren't explicitly documented.

docs/notebooks/layers_overview.ipynb

@@ -7,6 +7,15 @@
     "# 📙 `serket.nn` layers overview"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install git+https://github.com/ASEM000/serket --quiet"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -44,7 +53,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -126,7 +135,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -167,7 +176,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
@@ -176,14 +185,14 @@
        "Linear(\n",
        "  in_features=(10), \n",
        "  out_features=5, \n",
-       "  weight_init=he_normal, \n",
-       "  bias_init=ones, \n",
-       "  weight=f16[10,5](μ=0.08, σ=0.43, ∈[-1.01,0.87]), \n",
-       "  bias=f16[5](μ=1.00, σ=0.00, ∈[1.00,1.00])\n",
+       "  weight_init=glorot_uniform, \n",
+       "  bias_init=zeros, \n",
+       "  weight=f16[10,5](μ=0.07, σ=0.35, ∈[-0.63,0.60]), \n",
+       "  bias=f16[5](μ=0.00, σ=0.00, ∈[0.00,0.00])\n",
        ")"
       ]
      },
-     "execution_count": 3,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }

docs/notebooks/lazy_initialization.ipynb

@@ -15,14 +15,23 @@
    "cell_type": "code",
    "execution_count": 1,
    "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install git+https://github.com/ASEM000/serket --quiet"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
        "(5, 10)"
       ]
      },
-     "execution_count": 1,
+     "execution_count": 2,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -51,7 +60,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -60,7 +69,7 @@
        "(5, 10)"
       ]
      },
-     "execution_count": 2,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }

docs/notebooks/mental_model.ipynb

@@ -19,6 +19,15 @@
    "execution_count": 1,
    "metadata": {},
    "outputs": [],
+   "source": [
+    "!pip install git+https://github.com/ASEM000/serket --quiet"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "import jax\n",
     "import functools as ft\n",
@@ -62,7 +71,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {

docs/notebooks/regularization.ipynb

@@ -20,6 +20,15 @@
    "cell_type": "code",
    "execution_count": 1,
    "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install git+https://github.com/ASEM000/serket --quiet"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
@@ -67,7 +76,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {

docs/notebooks/train_bilstm.ipynb

@@ -7,6 +7,15 @@
     "# 🚆 Train bidirectional-`LSTM`"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install git+https://github.com/ASEM000/serket --quiet"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -16,7 +25,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -42,7 +51,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -70,7 +79,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -106,7 +115,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -149,23 +158,23 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch: 100/100\tBatch: 100/100\tBatch loss: 1.485198e-03\tTime: 0.018\r"
+      "Epoch: 100/100\tBatch: 100/100\tBatch loss: 1.485198e-03\tTime: 0.020\r"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "<matplotlib.legend.Legend at 0x16d857220>"
+       "<matplotlib.legend.Legend at 0x148c532b0>"
       ]
      },
-     "execution_count": 5,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     },