{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "This is a companion notebook for the book [Deep Learning with Python, Third Edition](https://www.manning.com/books/deep-learning-with-python-third-edition). For readability, it only contains runnable code blocks and section titles, and omits everything else in the book: text paragraphs, figures, and pseudocode.\n\n**If you want to be able to follow what's going on, I recommend reading the notebook side by side with your copy of the book.**\n\nThe book's contents are available online at [deeplearningwithpython.io](https://deeplearningwithpython.io)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "!pip install keras keras-hub --upgrade -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "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": {
    "colab_type": "text"
   },
   "source": [
    "## A deep dive on Keras"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### A spectrum of workflows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Different ways to build Keras models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### The Sequential model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import keras\n",
    "from keras import layers\n",
    "\n",
    "model = keras.Sequential(\n",
    "    [\n",
    "        layers.Dense(64, activation=\"relu\"),\n",
    "        layers.Dense(10, activation=\"softmax\"),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model = keras.Sequential()\n",
    "model.add(layers.Dense(64, activation=\"relu\"))\n",
    "model.add(layers.Dense(10, activation=\"softmax\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model.weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model.build(input_shape=(None, 3))\n",
    "model.weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model.summary(line_length=80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model = keras.Sequential(name=\"my_example_model\")\n",
    "model.add(layers.Dense(64, activation=\"relu\", name=\"my_first_layer\"))\n",
    "model.add(layers.Dense(10, activation=\"softmax\", name=\"my_last_layer\"))\n",
    "model.build((None, 3))\n",
    "model.summary(line_length=80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model = keras.Sequential()\n",
    "model.add(keras.Input(shape=(3,)))\n",
    "model.add(layers.Dense(64, activation=\"relu\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model.summary(line_length=80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model.add(layers.Dense(10, activation=\"softmax\"))\n",
    "model.summary(line_length=80)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### The Functional API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### A simple example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "inputs = keras.Input(shape=(3,), name=\"my_input\")\n",
    "features = layers.Dense(64, activation=\"relu\")(inputs)\n",
    "outputs = layers.Dense(10, activation=\"softmax\")(features)\n",
    "model = keras.Model(inputs=inputs, outputs=outputs, name=\"my_functional_model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "inputs = keras.Input(shape=(3,), name=\"my_input\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "inputs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "inputs.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "features = layers.Dense(64, activation=\"relu\")(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "outputs = layers.Dense(10, activation=\"softmax\")(features)\n",
    "model = keras.Model(inputs=inputs, outputs=outputs, name=\"my_functional_model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model.summary(line_length=80)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Multi-input, multi-output models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "vocabulary_size = 10000\n",
    "num_tags = 100\n",
    "num_departments = 4\n",
    "\n",
    "title = keras.Input(shape=(vocabulary_size,), name=\"title\")\n",
    "text_body = keras.Input(shape=(vocabulary_size,), name=\"text_body\")\n",
    "tags = keras.Input(shape=(num_tags,), name=\"tags\")\n",
    "\n",
    "features = layers.Concatenate()([title, text_body, tags])\n",
    "features = layers.Dense(64, activation=\"relu\", name=\"dense_features\")(features)\n",
    "\n",
    "priority = layers.Dense(1, activation=\"sigmoid\", name=\"priority\")(features)\n",
    "department = layers.Dense(\n",
    "    num_departments, activation=\"softmax\", name=\"department\"\n",
    ")(features)\n",
    "\n",
    "model = keras.Model(\n",
    "    inputs=[title, text_body, tags],\n",
    "    outputs=[priority, department],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Training a multi-input, multi-output model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "num_samples = 1280\n",
    "\n",
    "title_data = np.random.randint(0, 2, size=(num_samples, vocabulary_size))\n",
    "text_body_data = np.random.randint(0, 2, size=(num_samples, vocabulary_size))\n",
    "tags_data = np.random.randint(0, 2, size=(num_samples, num_tags))\n",
    "\n",
    "priority_data = np.random.random(size=(num_samples, 1))\n",
    "department_data = np.random.randint(0, num_departments, size=(num_samples, 1))\n",
    "\n",
    "model.compile(\n",
    "    optimizer=\"adam\",\n",
    "    loss=[\"mean_squared_error\", \"sparse_categorical_crossentropy\"],\n",
    "    metrics=[[\"mean_absolute_error\"], [\"accuracy\"]],\n",
    ")\n",
    "model.fit(\n",
    "    [title_data, text_body_data, tags_data],\n",
    "    [priority_data, department_data],\n",
    "    epochs=1,\n",
    ")\n",
    "model.evaluate(\n",
    "    [title_data, text_body_data, tags_data], [priority_data, department_data]\n",
    ")\n",
    "priority_preds, department_preds = model.predict(\n",
    "    [title_data, text_body_data, tags_data]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model.compile(\n",
    "    optimizer=\"adam\",\n",
    "    loss={\n",
    "        \"priority\": \"mean_squared_error\",\n",
    "        \"department\": \"sparse_categorical_crossentropy\",\n",
    "    },\n",
    "    metrics={\n",
    "        \"priority\": [\"mean_absolute_error\"],\n",
    "        \"department\": [\"accuracy\"],\n",
    "    },\n",
    ")\n",
    "model.fit(\n",
    "    {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data},\n",
    "    {\"priority\": priority_data, \"department\": department_data},\n",
    "    epochs=1,\n",
    ")\n",
    "model.evaluate(\n",
    "    {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data},\n",
    "    {\"priority\": priority_data, \"department\": department_data},\n",
    ")\n",
    "priority_preds, department_preds = model.predict(\n",
    "    {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### The power of the Functional API: Access to layer connectivity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "###### Plotting layer connectivity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "keras.utils.plot_model(model, \"ticket_classifier.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "keras.utils.plot_model(\n",
    "    model,\n",
    "    \"ticket_classifier_with_shape_info.png\",\n",
    "    show_shapes=True,\n",
    "    show_layer_names=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "###### Feature extraction with a Functional model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model.layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model.layers[3].input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model.layers[3].output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "features = model.layers[4].output\n",
    "difficulty = layers.Dense(3, activation=\"softmax\", name=\"difficulty\")(features)\n",
    "\n",
    "new_model = keras.Model(\n",
    "    inputs=[title, text_body, tags], outputs=[priority, department, difficulty]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "keras.utils.plot_model(\n",
    "    new_model,\n",
    "    \"updated_ticket_classifier.png\",\n",
    "    show_shapes=True,\n",
    "    show_layer_names=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Subclassing the Model class"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Rewriting our previous example as a subclassed model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "class CustomerTicketModel(keras.Model):\n",
    "    def __init__(self, num_departments):\n",
    "        super().__init__()\n",
    "        self.concat_layer = layers.Concatenate()\n",
    "        self.mixing_layer = layers.Dense(64, activation=\"relu\")\n",
    "        self.priority_scorer = layers.Dense(1, activation=\"sigmoid\")\n",
    "        self.department_classifier = layers.Dense(\n",
    "            num_departments, activation=\"softmax\"\n",
    "        )\n",
    "\n",
    "    def call(self, inputs):\n",
    "        title = inputs[\"title\"]\n",
    "        text_body = inputs[\"text_body\"]\n",
    "        tags = inputs[\"tags\"]\n",
    "\n",
    "        features = self.concat_layer([title, text_body, tags])\n",
    "        features = self.mixing_layer(features)\n",
    "        priority = self.priority_scorer(features)\n",
    "        department = self.department_classifier(features)\n",
    "        return priority, department"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model = CustomerTicketModel(num_departments=4)\n",
    "\n",
    "priority, department = model(\n",
    "    {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model.compile(\n",
    "    optimizer=\"adam\",\n",
    "    loss=[\"mean_squared_error\", \"sparse_categorical_crossentropy\"],\n",
    "    metrics=[[\"mean_absolute_error\"], [\"accuracy\"]],\n",
    ")\n",
    "model.fit(\n",
    "    {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data},\n",
    "    [priority_data, department_data],\n",
    "    epochs=1,\n",
    ")\n",
    "model.evaluate(\n",
    "    {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data},\n",
    "    [priority_data, department_data],\n",
    ")\n",
    "priority_preds, department_preds = model.predict(\n",
    "    {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Beware: What subclassed models don't support"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Mixing and matching different components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "class Classifier(keras.Model):\n",
    "    def __init__(self, num_classes=2):\n",
    "        super().__init__()\n",
    "        if num_classes == 2:\n",
    "            num_units = 1\n",
    "            activation = \"sigmoid\"\n",
    "        else:\n",
    "            num_units = num_classes\n",
    "            activation = \"softmax\"\n",
    "        self.dense = layers.Dense(num_units, activation=activation)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        return self.dense(inputs)\n",
    "\n",
    "inputs = keras.Input(shape=(3,))\n",
    "features = layers.Dense(64, activation=\"relu\")(inputs)\n",
    "outputs = Classifier(num_classes=10)(features)\n",
    "model = keras.Model(inputs=inputs, outputs=outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "inputs = keras.Input(shape=(64,))\n",
    "outputs = layers.Dense(1, activation=\"sigmoid\")(inputs)\n",
    "binary_classifier = keras.Model(inputs=inputs, outputs=outputs)\n",
    "\n",
    "class MyModel(keras.Model):\n",
    "    def __init__(self, num_classes=2):\n",
    "        super().__init__()\n",
    "        self.dense = layers.Dense(64, activation=\"relu\")\n",
    "        self.classifier = binary_classifier\n",
    "\n",
    "    def call(self, inputs):\n",
    "        features = self.dense(inputs)\n",
    "        return self.classifier(features)\n",
    "\n",
    "model = MyModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Remember: Use the right tool for the job"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Using built-in training and evaluation loops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "from keras.datasets import mnist\n",
    "\n",
    "def get_mnist_model():\n",
    "    inputs = keras.Input(shape=(28 * 28,))\n",
    "    features = layers.Dense(512, activation=\"relu\")(inputs)\n",
    "    features = layers.Dropout(0.5)(features)\n",
    "    outputs = layers.Dense(10, activation=\"softmax\")(features)\n",
    "    model = keras.Model(inputs, outputs)\n",
    "    return model\n",
    "\n",
    "(images, labels), (test_images, test_labels) = mnist.load_data()\n",
    "images = images.reshape((60000, 28 * 28)).astype(\"float32\") / 255\n",
    "test_images = test_images.reshape((10000, 28 * 28)).astype(\"float32\") / 255\n",
    "train_images, val_images = images[10000:], images[:10000]\n",
    "train_labels, val_labels = labels[10000:], labels[:10000]\n",
    "\n",
    "model = get_mnist_model()\n",
    "model.compile(\n",
    "    optimizer=\"adam\",\n",
    "    loss=\"sparse_categorical_crossentropy\",\n",
    "    metrics=[\"accuracy\"],\n",
    ")\n",
    "model.fit(\n",
    "    train_images,\n",
    "    train_labels,\n",
    "    epochs=3,\n",
    "    validation_data=(val_images, val_labels),\n",
    ")\n",
    "test_metrics = model.evaluate(test_images, test_labels)\n",
    "predictions = model.predict(test_images)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Writing your own metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "from keras import ops\n",
    "\n",
    "class RootMeanSquaredError(keras.metrics.Metric):\n",
    "    def __init__(self, name=\"rmse\", **kwargs):\n",
    "        super().__init__(name=name, **kwargs)\n",
    "        self.mse_sum = self.add_weight(name=\"mse_sum\", initializer=\"zeros\")\n",
    "        self.total_samples = self.add_weight(\n",
    "            name=\"total_samples\", initializer=\"zeros\"\n",
    "        )\n",
    "\n",
    "    def update_state(self, y_true, y_pred, sample_weight=None):\n",
    "        y_true = ops.one_hot(y_true, num_classes=ops.shape(y_pred)[1])\n",
    "        mse = ops.sum(ops.square(y_true - y_pred))\n",
    "        self.mse_sum.assign_add(mse)\n",
    "        num_samples = ops.shape(y_pred)[0]\n",
    "        self.total_samples.assign_add(num_samples)\n",
    "\n",
    "    def result(self):\n",
    "        return ops.sqrt(self.mse_sum / self.total_samples)\n",
    "\n",
    "    def reset_state(self):\n",
    "        self.mse_sum.assign(0.)\n",
    "        self.total_samples.assign(0.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model = get_mnist_model()\n",
    "model.compile(\n",
    "    optimizer=\"adam\",\n",
    "    loss=\"sparse_categorical_crossentropy\",\n",
    "    metrics=[\"accuracy\", RootMeanSquaredError()],\n",
    ")\n",
    "model.fit(\n",
    "    train_images,\n",
    "    train_labels,\n",
    "    epochs=3,\n",
    "    validation_data=(val_images, val_labels),\n",
    ")\n",
    "test_metrics = model.evaluate(test_images, test_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Using callbacks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### The EarlyStopping and ModelCheckpoint callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "callbacks_list = [\n",
    "    keras.callbacks.EarlyStopping(\n",
    "        monitor=\"accuracy\",\n",
    "        patience=1,\n",
    "    ),\n",
    "    keras.callbacks.ModelCheckpoint(\n",
    "        filepath=\"checkpoint_path.keras\",\n",
    "        monitor=\"val_loss\",\n",
    "        save_best_only=True,\n",
    "    ),\n",
    "]\n",
    "model = get_mnist_model()\n",
    "model.compile(\n",
    "    optimizer=\"adam\",\n",
    "    loss=\"sparse_categorical_crossentropy\",\n",
    "    metrics=[\"accuracy\"],\n",
    ")\n",
    "model.fit(\n",
    "    train_images,\n",
    "    train_labels,\n",
    "    epochs=10,\n",
    "    callbacks=callbacks_list,\n",
    "    validation_data=(val_images, val_labels),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model = keras.models.load_model(\"checkpoint_path.keras\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Writing your own callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "class LossHistory(keras.callbacks.Callback):\n",
    "    def on_train_begin(self, logs):\n",
    "        self.per_batch_losses = []\n",
    "\n",
    "    def on_batch_end(self, batch, logs):\n",
    "        self.per_batch_losses.append(logs.get(\"loss\"))\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs):\n",
    "        plt.clf()\n",
    "        plt.plot(\n",
    "            range(len(self.per_batch_losses)),\n",
    "            self.per_batch_losses,\n",
    "            label=\"Training loss for each batch\",\n",
    "        )\n",
    "        plt.xlabel(f\"Batch (epoch {epoch})\")\n",
    "        plt.ylabel(\"Loss\")\n",
    "        plt.legend()\n",
    "        plt.savefig(f\"plot_at_epoch_{epoch}\", dpi=300)\n",
    "        self.per_batch_losses = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model = get_mnist_model()\n",
    "model.compile(\n",
    "    optimizer=\"adam\",\n",
    "    loss=\"sparse_categorical_crossentropy\",\n",
    "    metrics=[\"accuracy\"],\n",
    ")\n",
    "model.fit(\n",
    "    train_images,\n",
    "    train_labels,\n",
    "    epochs=10,\n",
    "    callbacks=[LossHistory()],\n",
    "    validation_data=(val_images, val_labels),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Monitoring and visualization with TensorBoard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "model = get_mnist_model()\n",
    "model.compile(\n",
    "    optimizer=\"adam\",\n",
    "    loss=\"sparse_categorical_crossentropy\",\n",
    "    metrics=[\"accuracy\"],\n",
    ")\n",
    "\n",
    "tensorboard = keras.callbacks.TensorBoard(\n",
    "    log_dir=\"/full_path_to_your_log_dir\",\n",
    ")\n",
    "model.fit(\n",
    "    train_images,\n",
    "    train_labels,\n",
    "    epochs=10,\n",
    "    validation_data=(val_images, val_labels),\n",
    "    callbacks=[tensorboard],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%load_ext tensorboard\n",
    "%tensorboard --logdir /full_path_to_your_log_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Writing your own training and evaluation loops"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Training vs. inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Writing custom training step functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### A TensorFlow training step function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend tensorflow\n",
    "import tensorflow as tf\n",
    "\n",
    "model = get_mnist_model()\n",
    "loss_fn = keras.losses.SparseCategoricalCrossentropy()\n",
    "optimizer = keras.optimizers.Adam()\n",
    "\n",
    "def train_step(inputs, targets):\n",
    "    with tf.GradientTape() as tape:\n",
    "        predictions = model(inputs, training=True)\n",
    "        loss = loss_fn(targets, predictions)\n",
    "    gradients = tape.gradient(loss, model.trainable_weights)\n",
    "    optimizer.apply(gradients, model.trainable_weights)\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend tensorflow\n",
    "batch_size = 32\n",
    "inputs = train_images[:batch_size]\n",
    "targets = train_labels[:batch_size]\n",
    "loss = train_step(inputs, targets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### A PyTorch training step function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend torch\n",
    "import torch\n",
    "\n",
    "model = get_mnist_model()\n",
    "loss_fn = keras.losses.SparseCategoricalCrossentropy()\n",
    "optimizer = keras.optimizers.Adam()\n",
    "\n",
    "def train_step(inputs, targets):\n",
    "    predictions = model(inputs, training=True)\n",
    "    loss = loss_fn(targets, predictions)\n",
    "    loss.backward()\n",
    "    gradients = [weight.value.grad for weight in model.trainable_weights]\n",
    "    with torch.no_grad():\n",
    "        optimizer.apply(gradients, model.trainable_weights)\n",
    "    model.zero_grad()\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend torch\n",
    "batch_size = 32\n",
    "inputs = train_images[:batch_size]\n",
    "targets = train_labels[:batch_size]\n",
    "loss = train_step(inputs, targets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### A JAX training step function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend jax\n",
    "model = get_mnist_model()\n",
    "loss_fn = keras.losses.SparseCategoricalCrossentropy()\n",
    "\n",
    "def compute_loss_and_updates(\n",
    "    trainable_variables, non_trainable_variables, inputs, targets\n",
    "):\n",
    "    outputs, non_trainable_variables = model.stateless_call(\n",
    "        trainable_variables, non_trainable_variables, inputs, training=True\n",
    "    )\n",
    "    loss = loss_fn(targets, outputs)\n",
    "    return loss, non_trainable_variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend jax\n",
    "import jax\n",
    "\n",
    "grad_fn = jax.value_and_grad(compute_loss_and_updates, has_aux=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend jax\n",
    "optimizer = keras.optimizers.Adam()\n",
    "optimizer.build(model.trainable_variables)\n",
    "\n",
    "def train_step(state, inputs, targets):\n",
    "    (trainable_variables, non_trainable_variables, optimizer_variables) = state\n",
    "    (loss, non_trainable_variables), grads = grad_fn(\n",
    "        trainable_variables, non_trainable_variables, inputs, targets\n",
    "    )\n",
    "    trainable_variables, optimizer_variables = optimizer.stateless_apply(\n",
    "        optimizer_variables, grads, trainable_variables\n",
    "    )\n",
    "    return loss, (\n",
    "        trainable_variables,\n",
    "        non_trainable_variables,\n",
    "        optimizer_variables,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend jax\n",
    "batch_size = 32\n",
    "inputs = train_images[:batch_size]\n",
    "targets = train_labels[:batch_size]\n",
    "\n",
    "trainable_variables = [v.value for v in model.trainable_variables]\n",
    "non_trainable_variables = [v.value for v in model.non_trainable_variables]\n",
    "optimizer_variables = [v.value for v in optimizer.variables]\n",
    "\n",
    "state = (trainable_variables, non_trainable_variables, optimizer_variables)\n",
    "loss, state = train_step(state, inputs, targets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Low-level usage of metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "from keras import ops\n",
    "\n",
    "metric = keras.metrics.SparseCategoricalAccuracy()\n",
    "targets = ops.array([0, 1, 2])\n",
    "predictions = ops.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])\n",
    "metric.update_state(targets, predictions)\n",
    "current_result = metric.result()\n",
    "print(f\"result: {current_result:.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "values = ops.array([0, 1, 2, 3, 4])\n",
    "mean_tracker = keras.metrics.Mean()\n",
    "for value in values:\n",
    "    mean_tracker.update_state(value)\n",
    "print(f\"Mean of values: {mean_tracker.result():.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "metric = keras.metrics.SparseCategoricalAccuracy()\n",
    "targets = ops.array([0, 1, 2])\n",
    "predictions = ops.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])\n",
    "\n",
    "metric_variables = metric.variables\n",
    "metric_variables = metric.stateless_update_state(\n",
    "    metric_variables, targets, predictions\n",
    ")\n",
    "current_result = metric.stateless_result(metric_variables)\n",
    "print(f\"result: {current_result:.2f}\")\n",
    "\n",
    "metric_variables = metric.stateless_reset_state()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Using fit() with a custom training loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Customizing fit() with TensorFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend tensorflow\n",
    "import keras\n",
    "from keras import layers\n",
    "\n",
    "loss_fn = keras.losses.SparseCategoricalCrossentropy()\n",
    "loss_tracker = keras.metrics.Mean(name=\"loss\")\n",
    "\n",
    "class CustomModel(keras.Model):\n",
    "    def train_step(self, data):\n",
    "        inputs, targets = data\n",
    "        with tf.GradientTape() as tape:\n",
    "            predictions = self(inputs, training=True)\n",
    "            loss = loss_fn(targets, predictions)\n",
    "        gradients = tape.gradient(loss, self.trainable_weights)\n",
    "        self.optimizer.apply(gradients, self.trainable_weights)\n",
    "\n",
    "        loss_tracker.update_state(loss)\n",
    "        return {\"loss\": loss_tracker.result()}\n",
    "\n",
    "    @property\n",
    "    def metrics(self):\n",
    "        return [loss_tracker]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend tensorflow\n",
    "def get_custom_model():\n",
    "    inputs = keras.Input(shape=(28 * 28,))\n",
    "    features = layers.Dense(512, activation=\"relu\")(inputs)\n",
    "    features = layers.Dropout(0.5)(features)\n",
    "    outputs = layers.Dense(10, activation=\"softmax\")(features)\n",
    "    model = CustomModel(inputs, outputs)\n",
    "    model.compile(optimizer=keras.optimizers.Adam())\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend tensorflow\n",
    "model = get_custom_model()\n",
    "model.fit(train_images, train_labels, epochs=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Customizing fit() with PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend torch\n",
    "import keras\n",
    "from keras import layers\n",
    "\n",
    "loss_fn = keras.losses.SparseCategoricalCrossentropy()\n",
    "loss_tracker = keras.metrics.Mean(name=\"loss\")\n",
    "\n",
    "class CustomModel(keras.Model):\n",
    "    def train_step(self, data):\n",
    "        inputs, targets = data\n",
    "        predictions = self(inputs, training=True)\n",
    "        loss = loss_fn(targets, predictions)\n",
    "\n",
    "        loss.backward()\n",
    "        trainable_weights = [v for v in self.trainable_weights]\n",
    "        gradients = [v.value.grad for v in trainable_weights]\n",
    "\n",
    "        with torch.no_grad():\n",
    "            self.optimizer.apply(gradients, trainable_weights)\n",
    "\n",
    "        loss_tracker.update_state(loss)\n",
    "        return {\"loss\": loss_tracker.result()}\n",
    "\n",
    "    @property\n",
    "    def metrics(self):\n",
    "        return [loss_tracker]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend torch\n",
    "def get_custom_model():\n",
    "    inputs = keras.Input(shape=(28 * 28,))\n",
    "    features = layers.Dense(512, activation=\"relu\")(inputs)\n",
    "    features = layers.Dropout(0.5)(features)\n",
    "    outputs = layers.Dense(10, activation=\"softmax\")(features)\n",
    "    model = CustomModel(inputs, outputs)\n",
    "    model.compile(optimizer=keras.optimizers.Adam())\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend torch\n",
    "model = get_custom_model()\n",
    "model.fit(train_images, train_labels, epochs=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Customizing fit() with JAX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend jax\n",
    "import keras\n",
    "from keras import layers\n",
    "\n",
    "loss_fn = keras.losses.SparseCategoricalCrossentropy()\n",
    "\n",
    "class CustomModel(keras.Model):\n",
    "    def compute_loss_and_updates(\n",
    "        self,\n",
    "        trainable_variables,\n",
    "        non_trainable_variables,\n",
    "        inputs,\n",
    "        targets,\n",
    "        training=False,\n",
    "    ):\n",
    "        predictions, non_trainable_variables = self.stateless_call(\n",
    "            trainable_variables,\n",
    "            non_trainable_variables,\n",
    "            inputs,\n",
    "            training=training,\n",
    "        )\n",
    "        loss = loss_fn(targets, predictions)\n",
    "        return loss, non_trainable_variables\n",
    "\n",
    "    def train_step(self, state, data):\n",
    "        (\n",
    "            trainable_variables,\n",
    "            non_trainable_variables,\n",
    "            optimizer_variables,\n",
    "            metrics_variables,\n",
    "        ) = state\n",
    "        inputs, targets = data\n",
    "\n",
    "        grad_fn = jax.value_and_grad(\n",
    "            self.compute_loss_and_updates, has_aux=True\n",
    "        )\n",
    "\n",
    "        (loss, non_trainable_variables), grads = grad_fn(\n",
    "            trainable_variables,\n",
    "            non_trainable_variables,\n",
    "            inputs,\n",
    "            targets,\n",
    "            training=True,\n",
    "        )\n",
    "\n",
    "        (\n",
    "            trainable_variables,\n",
    "            optimizer_variables,\n",
    "        ) = self.optimizer.stateless_apply(\n",
    "            optimizer_variables, grads, trainable_variables\n",
    "        )\n",
    "\n",
    "        logs = {\"loss\": loss}\n",
    "        state = (\n",
    "            trainable_variables,\n",
    "            non_trainable_variables,\n",
    "            optimizer_variables,\n",
    "            metrics_variables,\n",
    "        )\n",
    "        return logs, state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend jax\n",
    "def get_custom_model():\n",
    "    inputs = keras.Input(shape=(28 * 28,))\n",
    "    features = layers.Dense(512, activation=\"relu\")(inputs)\n",
    "    features = layers.Dropout(0.5)(features)\n",
    "    outputs = layers.Dense(10, activation=\"softmax\")(features)\n",
    "    model = CustomModel(inputs, outputs)\n",
    "    model.compile(optimizer=keras.optimizers.Adam())\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend jax\n",
    "model = get_custom_model()\n",
    "model.fit(train_images, train_labels, epochs=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Handling metrics in a custom train_step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### train_step() metrics handling with TensorFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend tensorflow\n",
    "import keras\n",
    "from keras import layers\n",
    "\n",
    "class CustomModel(keras.Model):\n",
    "    def train_step(self, data):\n",
    "        inputs, targets = data\n",
    "        with tf.GradientTape() as tape:\n",
    "            predictions = self(inputs, training=True)\n",
    "            loss = self.compute_loss(y=targets, y_pred=predictions)\n",
    "\n",
    "        gradients = tape.gradient(loss, self.trainable_weights)\n",
    "        self.optimizer.apply(gradients, self.trainable_weights)\n",
    "\n",
    "        for metric in self.metrics:\n",
    "            if metric.name == \"loss\":\n",
    "                metric.update_state(loss)\n",
    "            else:\n",
    "                metric.update_state(targets, predictions)\n",
    "\n",
    "        return {m.name: m.result() for m in self.metrics}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend tensorflow\n",
    "def get_custom_model():\n",
    "    inputs = keras.Input(shape=(28 * 28,))\n",
    "    features = layers.Dense(512, activation=\"relu\")(inputs)\n",
    "    features = layers.Dropout(0.5)(features)\n",
    "    outputs = layers.Dense(10, activation=\"softmax\")(features)\n",
    "    model = CustomModel(inputs, outputs)\n",
    "    model.compile(\n",
    "        optimizer=keras.optimizers.Adam(),\n",
    "        loss=keras.losses.SparseCategoricalCrossentropy(),\n",
    "        metrics=[keras.metrics.SparseCategoricalAccuracy()],\n",
    "    )\n",
    "    return model\n",
    "\n",
    "model = get_custom_model()\n",
    "model.fit(train_images, train_labels, epochs=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### train_step() metrics handling with PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend torch\n",
    "import keras\n",
    "from keras import layers\n",
    "\n",
    "class CustomModel(keras.Model):\n",
    "    def train_step(self, data):\n",
    "        inputs, targets = data\n",
    "        predictions = self(inputs, training=True)\n",
    "        loss = self.compute_loss(y=targets, y_pred=predictions)\n",
    "\n",
    "        loss.backward()\n",
    "        trainable_weights = [v for v in self.trainable_weights]\n",
    "        gradients = [v.value.grad for v in trainable_weights]\n",
    "\n",
    "        with torch.no_grad():\n",
    "            self.optimizer.apply(gradients, trainable_weights)\n",
    "\n",
    "        for metric in self.metrics:\n",
    "            if metric.name == \"loss\":\n",
    "                metric.update_state(loss)\n",
    "            else:\n",
    "                metric.update_state(targets, predictions)\n",
    "\n",
    "        return {m.name: m.result() for m in self.metrics}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend torch\n",
    "def get_custom_model():\n",
    "    inputs = keras.Input(shape=(28 * 28,))\n",
    "    features = layers.Dense(512, activation=\"relu\")(inputs)\n",
    "    features = layers.Dropout(0.5)(features)\n",
    "    outputs = layers.Dense(10, activation=\"softmax\")(features)\n",
    "    model = CustomModel(inputs, outputs)\n",
    "    model.compile(\n",
    "        optimizer=keras.optimizers.Adam(),\n",
    "        loss=keras.losses.SparseCategoricalCrossentropy(),\n",
    "        metrics=[keras.metrics.SparseCategoricalAccuracy()],\n",
    "    )\n",
    "    return model\n",
    "\n",
    "model = get_custom_model()\n",
    "model.fit(train_images, train_labels, epochs=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### train_step() metrics handling with JAX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "%%backend jax\n",
    "import keras\n",
    "from keras import layers\n",
    "\n",
    "class CustomModel(keras.Model):\n",
    "    def compute_loss_and_updates(\n",
    "        self,\n",
    "        trainable_variables,\n",
    "        non_trainable_variables,\n",
    "        inputs,\n",
    "        targets,\n",
    "        training=False,\n",
    "    ):\n",
    "        predictions, non_trainable_variables = self.stateless_call(\n",
    "            trainable_variables,\n",
    "            non_trainable_variables,\n",
    "            inputs,\n",
    "            training=training,\n",
    "        )\n",
    "        loss = self.compute_loss(y=targets, y_pred=predictions)\n",
    "        return loss, (predictions, non_trainable_variables)\n",
    "\n",
    "    def train_step(self, state, data):\n",
    "        (\n",
    "            trainable_variables,\n",
    "            non_trainable_variables,\n",
    "            optimizer_variables,\n",
    "            metrics_variables,\n",
    "        ) = state\n",
    "        inputs, targets = data\n",
    "\n",
    "        grad_fn = jax.value_and_grad(\n",
    "            self.compute_loss_and_updates, has_aux=True\n",
    "        )\n",
    "\n",
    "        (loss, (predictions, non_trainable_variables)), grads = grad_fn(\n",
    "            trainable_variables,\n",
    "            non_trainable_variables,\n",
    "            inputs,\n",
    "            targets,\n",
    "            training=True,\n",
    "        )\n",
    "        (\n",
    "            trainable_variables,\n",
    "            optimizer_variables,\n",
    "        ) = self.optimizer.stateless_apply(\n",
    "            optimizer_variables, grads, trainable_variables\n",
    "        )\n",
    "\n",
    "        new_metrics_vars = []\n",
    "        logs = {}\n",
    "        for metric in self.metrics:\n",
    "            num_prev = len(new_metrics_vars)\n",
    "            num_current = len(metric.variables)\n",
    "            current_vars = metrics_variables[num_prev : num_prev + num_current]\n",
    "            if metric.name == \"loss\":\n",
    "                current_vars = metric.stateless_update_state(current_vars, loss)\n",
    "            else:\n",
    "                current_vars = metric.stateless_update_state(\n",
    "                    current_vars, targets, predictions\n",
    "                )\n",
    "            logs[metric.name] = metric.stateless_result(current_vars)\n",
    "            new_metrics_vars += current_vars\n",
    "\n",
    "        state = (\n",
    "            trainable_variables,\n",
    "            non_trainable_variables,\n",
    "            optimizer_variables,\n",
    "            new_metrics_vars,\n",
    "        )\n",
    "        return logs, state"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "chapter07_deep-dive-keras",
   "private_outputs": false,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}