Remove set_backend() from all notebooks

Author: mattdangerw

chapter02_mathematical-building-blocks.ipynb

@@ -32,6 +32,31 @@
     "os.environ[\"KERAS_BACKEND\"] = \"jax\""
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "cellView": "form",
+    "colab_type": "code"
+   },
+   "outputs": [],
+   "source": [
+    "# @title\n",
+    "import os\n",
+    "from IPython.core.magic import register_cell_magic\n",
+    "\n",
+    "@register_cell_magic\n",
+    "def backend(line, cell):\n",
+    "    current, required = os.environ.get(\"KERAS_BACKEND\", \"\"), line.split()[-1]\n",
+    "    if current == required:\n",
+    "        get_ipython().run_cell(cell)\n",
+    "    else:\n",
+    "        print(\n",
+    "            f\"This cell requires the {required} backend. To run it, change KERAS_BACKEND to \"\n",
+    "            f\"\\\"{required}\\\" at the top of the notebook, restart the runtime, and rerun the notebook.\"\n",
+    "        )"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -1113,17 +1138,6 @@
     "##### A simple Dense class"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 0,
-   "metadata": {
-    "colab_type": "code"
-   },
-   "outputs": [],
-   "source": [
-    "keras.config.set_backend(\"tensorflow\")"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 0,
@@ -1308,6 +1322,7 @@
    },
    "outputs": [],
    "source": [
+    "%%backend tensorflow\n",
     "import tensorflow as tf\n",
     "\n",
     "x = tf.zeros(shape=())\n",
@@ -1324,6 +1339,7 @@
    },
    "outputs": [],
    "source": [
+    "%%backend tensorflow\n",
     "def one_training_step(model, images_batch, labels_batch):\n",
     "    with tf.GradientTape() as tape:\n",
     "        predictions = model(images_batch)\n",
@@ -1351,6 +1367,7 @@
    },
    "outputs": [],
    "source": [
+    "%%backend tensorflow\n",
     "def fit(model, images, labels, epochs, batch_size=128):\n",
     "    for epoch_counter in range(epochs):\n",
     "        print(f\"Epoch {epoch_counter}\")\n",
@@ -1370,6 +1387,7 @@
    },
    "outputs": [],
    "source": [
+    "%%backend tensorflow\n",
     "from keras.datasets import mnist\n",
     "\n",
     "(train_images, train_labels), (test_images, test_labels) = mnist.load_data()\n",
@@ -1399,6 +1417,7 @@
    },
    "outputs": [],
    "source": [
+    "%%backend tensorflow\n",
     "predictions = model(test_images)\n",
     "predicted_labels = ops.argmax(predictions, axis=1)\n",
     "matches = predicted_labels == test_labels\n",

chapter03_introduction-to-ml-frameworks.ipynb

@@ -32,6 +32,31 @@
     "os.environ[\"KERAS_BACKEND\"] = \"jax\""
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "cellView": "form",
+    "colab_type": "code"
+   },
+   "outputs": [],
+   "source": [
+    "# @title\n",
+    "import os\n",
+    "from IPython.core.magic import register_cell_magic\n",
+    "\n",
+    "@register_cell_magic\n",
+    "def backend(line, cell):\n",
+    "    current, required = os.environ.get(\"KERAS_BACKEND\", \"\"), line.split()[-1]\n",
+    "    if current == required:\n",
+    "        get_ipython().run_cell(cell)\n",
+    "    else:\n",
+    "        print(\n",
+    "            f\"This cell requires the {required} backend. To run it, change KERAS_BACKEND to \"\n",
+    "            f\"\\\"{required}\\\" at the top of the notebook, restart the runtime, and rerun the notebook.\"\n",
+    "        )"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {

chapter04_classification-and-regression.ipynb

@@ -32,6 +32,31 @@
     "os.environ[\"KERAS_BACKEND\"] = \"jax\""
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "cellView": "form",
+    "colab_type": "code"
+   },
+   "outputs": [],
+   "source": [
+    "# @title\n",
+    "import os\n",
+    "from IPython.core.magic import register_cell_magic\n",
+    "\n",
+    "@register_cell_magic\n",
+    "def backend(line, cell):\n",
+    "    current, required = os.environ.get(\"KERAS_BACKEND\", \"\"), line.split()[-1]\n",
+    "    if current == required:\n",
+    "        get_ipython().run_cell(cell)\n",
+    "    else:\n",
+    "        print(\n",
+    "            f\"This cell requires the {required} backend. To run it, change KERAS_BACKEND to \"\n",
+    "            f\"\\\"{required}\\\" at the top of the notebook, restart the runtime, and rerun the notebook.\"\n",
+    "        )"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {

chapter05_fundamentals-of-ml.ipynb

@@ -32,6 +32,31 @@
     "os.environ[\"KERAS_BACKEND\"] = \"jax\""
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "cellView": "form",
+    "colab_type": "code"
+   },
+   "outputs": [],
+   "source": [
+    "# @title\n",
+    "import os\n",
+    "from IPython.core.magic import register_cell_magic\n",
+    "\n",
+    "@register_cell_magic\n",
+    "def backend(line, cell):\n",
+    "    current, required = os.environ.get(\"KERAS_BACKEND\", \"\"), line.split()[-1]\n",
+    "    if current == required:\n",
+    "        get_ipython().run_cell(cell)\n",
+    "    else:\n",
+    "        print(\n",
+    "            f\"This cell requires the {required} backend. To run it, change KERAS_BACKEND to \"\n",
+    "            f\"\\\"{required}\\\" at the top of the notebook, restart the runtime, and rerun the notebook.\"\n",
+    "        )"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {