{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "mnist_tensorflow.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "PZtRtMMUZHJS"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PZtRtMMUZHJS",
        "colab_type": "text"
      },
      "source": [
        "##### Copyright 2020 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TouZL3JZZSQe",
        "colab_type": "code",
        "colab": {},
        "cellView": "both"
      },
      "source": [
        "#@title License header\n",
        "# Copyright 2020 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#      https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O6c3qfq5Zv57",
        "colab_type": "text"
      },
      "source": [
        "# MNIST Model TensorFlow Training, IREE Execution\n",
        "\n",
        "## Overview\n",
        "\n",
        "This notebook creates and trains a TensorFlow 2.0 model for recognizing handwritten digits using the [MNIST dataset](https://en.wikipedia.org/wiki/MNIST_database), then compiles and executes that trained model using IREE.\n",
        "\n",
        "## Running Locally\n",
        "\n",
        "*  Refer to [iree/docs/using_colab.md](https://github.com/google/iree/blob/master/docs/using_colab.md) for general information\n",
        "*  Ensure that you have a recent version of TensorFlow 2.0 [installed on your system](https://www.tensorflow.org/install)\n",
        "*  Enable IREE/TF integration by adding to your user.bazelrc: `build --define=iree_tensorflow=true`\n",
        "*  Start colab by running `python colab/start_colab_kernel.py` (see that file for additional instructions)\n",
        "*  Note: you may need to restart your runtime in order to re-run certain cells. Some of the APIs are not yet stable enough for repeated invocations"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wBXlE69Ia2QU",
        "colab_type": "text"
      },
      "source": [
        "# Setup Steps"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EPF7RGQDYK-M",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from matplotlib import pyplot as plt\n",
        "from pyiree.tf import compiler as ireec\n",
        "from pyiree import rt as ireert\n",
        "\n",
        "tf.compat.v1.enable_eager_execution()\n",
        "\n",
        "SAVE_PATH = os.path.join(os.environ[\"HOME\"], \"saved_models\")\n",
        "os.makedirs(SAVE_PATH, exist_ok=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "43BH_9YcsGs8",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "cellView": "form",
        "outputId": "e24a17ea-297b-491b-f97c-244b411ea8d1"
      },
      "source": [
        "#@title Notebook settings { run: \"auto\" }\n",
        "\n",
        "#@markdown -----\n",
        "#@markdown ### Configuration\n",
        "\n",
        "backend_choice = \"GPU (vulkan-spirv)\" #@param [ \"GPU (vulkan-spirv)\", \"CPU (interpreter bytecode)\" ]\n",
        "\n",
        "if backend_choice == \"GPU (vulkan-spirv)\":\n",
        "  backend_name = \"vulkan-spirv\"\n",
        "  driver_name = \"vulkan\"\n",
        "else:\n",
        "  backend_name = \"vmla\"\n",
        "  driver_name = \"vmla\"\n",
        "tf.print(\"Using IREE compiler backend '%s' and runtime driver '%s'\" % (backend_name, driver_name))\n",
        "\n",
        "#@markdown -----\n",
        "#@markdown ### Training Parameters\n",
        "\n",
        "#@markdown <sup>Batch size used to subdivide the training and evaluation samples</sup>\n",
        "batch_size = 200  #@param { type: \"slider\", min: 10, max: 400 }\n",
        "\n",
        "#@markdown <sup>Epochs for training/eval. Higher values take longer to run but generally produce more accurate models</sup>\n",
        "num_epochs = 5    #@param { type: \"slider\", min:  1, max:  20 }\n",
        "\n",
        "#@markdown -----"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Using IREE compiler backend 'vulkan-spirv' and runtime driver 'vulkan'\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5vkQOMOMbXdy",
        "colab_type": "text"
      },
      "source": [
        "# Create and Train MNIST Model in TensorFlow\n",
        "\n",
        "The specific details of the training process here aren't critical to the model compilation and execution through IREE."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GXZIrReTbTHN",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 486
        },
        "cellView": "form",
        "outputId": "d727261b-7457-47ba-97c6-aca339a9d91f"
      },
      "source": [
        "#@title Load MNIST dataset, setup training and evaluation\n",
        "\n",
        "NUM_CLASSES = 10  # One per digit [0, 1, 2, ..., 9]\n",
        "IMG_ROWS, IMG_COLS = 28, 28\n",
        "\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "tf.print(\"Loaded MNIST dataset!\")\n",
        "\n",
        "x_train = x_train.reshape(x_train.shape[0], IMG_ROWS, IMG_COLS, 1)\n",
        "x_test = x_test.reshape(x_test.shape[0], IMG_ROWS, IMG_COLS, 1)\n",
        "input_shape = (IMG_ROWS, IMG_COLS, 1)\n",
        "\n",
        "# Scale pixel values from [0, 255] integers to [0.0, 1.0] floats.\n",
        "x_train = x_train.astype(\"float32\") / 255\n",
        "x_test = x_test.astype(\"float32\") / 255\n",
        "\n",
        "steps_per_epoch = int(x_train.shape[0] / batch_size)\n",
        "steps_per_eval = int(x_test.shape[0] / batch_size)\n",
        "\n",
        "# Convert class vectors to binary class matrices.\n",
        "y_train = tf.keras.utils.to_categorical(y_train, NUM_CLASSES)\n",
        "y_test = tf.keras.utils.to_categorical(y_test, NUM_CLASSES)\n",
        "\n",
        "# Construct batched datasets for training/evaluation.\n",
        "train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_dataset = train_dataset.batch(batch_size, drop_remainder=True)\n",
        "test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
        "test_dataset = test_dataset.batch(batch_size, drop_remainder=True)\n",
        "\n",
        "# Create a distribution strategy for the dataset (single machine).\n",
        "strategy = tf.distribute.experimental.CentralStorageStrategy()\n",
        "train_dist_ds = strategy.experimental_distribute_dataset(train_dataset)\n",
        "test_dist_ds = strategy.experimental_distribute_dataset(test_dataset)\n",
        "\n",
        "tf.print(\"Configured data for training and evaluation!\")\n",
        "tf.print(\"  sample shape: %s\" % str(x_train[0].shape))\n",
        "tf.print(\"  training samples: %s\" % x_train.shape[0])\n",
        "tf.print(\"  test     samples: %s\" % x_test.shape[0])\n",
        "tf.print(\"  epochs: %s\" % num_epochs)\n",
        "tf.print(\"  steps/epoch: %s\" % steps_per_epoch)\n",
        "tf.print(\"  steps/eval : %s\" % steps_per_eval)\n",
        "\n",
        "tf.print(\"\")\n",
        "tf.print(\"Sample image from the dataset:\")\n",
        "SAMPLE_EXAMPLE_INDEX = 1\n",
        "sample_image = x_test[SAMPLE_EXAMPLE_INDEX]\n",
        "sample_image_batch = np.expand_dims(sample_image, axis=0)\n",
        "sample_label = y_test[SAMPLE_EXAMPLE_INDEX]\n",
        "plt.imshow(sample_image.reshape(IMG_ROWS, IMG_COLS))\n",
        "plt.show()\n",
        "tf.print(\"\\nGround truth labels: %s\" % str(sample_label))"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loaded MNIST dataset!\n",
            "INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ('/device:CPU:0',), variable_device = '/device:CPU:0'\n",
            "Configured data for training and evaluation!\n",
            "  sample shape: (28, 28, 1)\n",
            "  training samples: 60000\n",
            "  test     samples: 10000\n",
            "  epochs: 5\n",
            "  steps/epoch: 300\n",
            "  steps/eval : 50\n",
            "\n",
            "Sample image from the dataset:\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD4CAYAAAAq5pAIAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy8li6FKAAANzUlEQVR4nO3df6zV9X3H8dcL5IdFVBiMMSRaLMRiF6G9oXV1m8a1s/xRbLK5ks5hY3O7rG5tQtIat6Q2/RGzVN2WNV1oJaWLP+L8UVlqOpHaOFuCXhwFhLZQhyvsChJuB24ZcK/v/XG/NFe93++5nPM9P+T9fCQ355zv+3y/33eOvvie8/2c7/k4IgTg7Dep2w0A6AzCDiRB2IEkCDuQBGEHkjinkzub6mkxXTM6uUsglf/T/+hknPB4tZbCbvs6SX8nabKkb0bEHVXPn64Zeq+vbWWXACpsjc2ltabfxtueLOlrkj4kaamk1baXNrs9AO3Vymf2FZL2RcSLEXFS0gOSVtXTFoC6tRL2BZJ+MebxgWLZ69jutz1ge+CUTrSwOwCtaPvZ+IhYFxF9EdE3RdPavTsAJVoJ+0FJC8c8vqhYBqAHtRL25yQttv1221MlfVTSxnraAlC3pofeImLY9i2S/lWjQ2/rI+KF2joDUKuWxtkj4nFJj9fUC4A24uuyQBKEHUiCsANJEHYgCcIOJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJwg4k0dGfkkZz9n/pysr6yPTyyTnnXv5K5bpbrni4qZ5Ou/T7H6+sz3z23NLavL//UUv7xpnhyA4kQdiBJAg7kARhB5Ig7EAShB1IgrADSTDO3gOGvru4sr5r2T+0bd+nyofoJ+Qn13yzsn5v3/zS2oObfq9y3ZE9e5vqCePjyA4kQdiBJAg7kARhB5Ig7EAShB1IgrADSTDO3gGNxtF/uOyBtu37H3+5qLJ+15YPVNYvubj6evgnlj5SWf/YzMHS2pdvmlO57qLPMc5ep5bCbnu/pOOSRiQNR0RfHU0BqF8dR/ZrIuJIDdsB0EZ8ZgeSaDXsIekJ29ts94/3BNv9tgdsD5zSiRZ3B6BZrb6NvyoiDtr+dUmbbP8kIp4e+4SIWCdpnSSd79ktXnYBoFktHdkj4mBxe1jSo5JW1NEUgPo1HXbbM2zPPH1f0gcl7aqrMQD1auVt/DxJj9o+vZ37IuJ7tXT1FjN87Xsq69+/4msNtjClsvq3Q0sq60/9ccWI538drlx3ydBAZX3S9OmV9a9s/a3K+m1zdpbWhmcNV66LejUd9oh4UdIVNfYCoI0YegOSIOxAEoQdSIKwA0kQdiAJLnGtwasLplbWJzX4N7XR0NoPPlw9vDXy4k8r663Y94XllfX7Zt/ZYAvTSisXfY9jTSfxagNJEHYgCcIOJEHYgSQIO5AEYQeSIOxAEoyz1+DCb2+prP/hwJ9U1j10rLI+PLj/DDuqzydWPllZP29S+Tg6egtHdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgnH2DhjZ/bNut1Bq/5evrKzffOFXG2yh+qem1w6+r7Q288k9leuONNgzzgxHdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgnH2s9wvb6weR//hn1aPo18wqXocfcuJyZX17V8q/935c489W7ku6tXwyG57ve3DtneNWTbb9ibbe4vbWe1tE0CrJvI2/luSrnvDslslbY6IxZI2F48B9LCGYY+IpyUdfcPiVZI2FPc3SLq+5r4A1KzZz+zzImKwuP+ypHllT7TdL6lfkqbrbU3uDkCrWj4bHxEhKSrq6yKiLyL6plRM8gegvZoN+yHb8yWpuD1cX0sA2qHZsG+UtKa4v0bSY/W0A6BdGn5mt32/pKslzbF9QNLnJd0h6UHbN0t6SdIN7WwSzTvy7tJPWJIaj6M3suYHn6isL/kOY+m9omHYI2J1SenamnsB0EZ8XRZIgrADSRB2IAnCDiRB2IEkuMT1LHBy08WltS2X3dlg7eqhtyu2rKmsv3Ptzyvr/Bx07+DIDiRB2IEkCDuQBGEHkiDsQBKEHUiCsANJMM7+FnDOoksq6198xz+X1mY1uIR124nqfV/8xeqR8pGhoeoNoGdwZAeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJBhnfwu49MGDlfXlU5v/N3v15j+rrC/58XNNbxu9hSM7kARhB5Ig7EAShB1IgrADSRB2IAnCDiTBOHsPGFpzZWX9C/Ma/fb7tNLKmv2/X7nmOz+7r7LO776fPRoe2W2vt33Y9q4xy263fdD29uJvZXvbBNCqibyN/5ak68ZZfndELCv+Hq+3LQB1axj2iHha0tEO9AKgjVo5QXeL7R3F2/xZZU+y3W97wPbAKTX4wTMAbdNs2L8u6VJJyyQNSio9gxQR6yKiLyL6plScSALQXk2FPSIORcRIRLwm6RuSVtTbFoC6NRV22/PHPPyIpF1lzwXQGxqOs9u+X9LVkubYPiDp85Kutr1MUkjaL+mTbezxLe+cBb9ZWf+dv9xaWT9vUvMff7bsfkdlfckQ16tn0TDsEbF6nMX3tKEXAG3E12WBJAg7kARhB5Ig7EAShB1IgktcO2DPbQsr69/5jX9pafvX7Pyj0hqXsOI0juxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kATj7B2w7cN3N3hGa7/gc8Gfv1ZaGx4aamnbOHtwZAeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJBhnPwucmndBaW3KyQUd7OTNRl45UlqLE9XTgXla9fcPJs+d01RPkjQy98LK+t61U5ve9kTEiEtrl/1Fg98gOHasqX1yZAeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJBhnPwt896H13W6h1G//+3iTAI86cuj8ynVnzT1eWd/6nvua6qnXLf3rWyrriz67pantNjyy215o+ynbu22/YPvTxfLZtjfZ3lvczmqqAwAdMZG38cOS1kbEUknvk/Qp20sl3Sppc0QslrS5eAygRzUMe0QMRsTzxf3jkvZIWiBplaQNxdM2SLq+XU0CaN0ZfWa3fYmk5ZK2SpoXEYNF6WVJ80rW6ZfUL0nT9bZm+wTQogmfjbd9nqSHJX0mIl73TfyICEkx3noRsS4i+iKib0qLP6wIoHkTCrvtKRoN+r0R8Uix+JDt+UV9vqTD7WkRQB0avo23bUn3SNoTEXeNKW2UtEbSHcXtY23p8CywavfHKuub3/VQhzrpvB8tv79r+/7fOFlaOxXlP789ESt33FRZ/+/tzV9+u+CZ4abXrTKRz+zvl3SjpJ22txfLbtNoyB+0fbOklyTd0JYOAdSiYdgj4hlJZVfaX1tvOwDaha/LAkkQdiAJwg4kQdiBJAg7kASXuHbAuX/wH5X1y79SfUljtPG/0szLjlbW23kZ6eX/9vHKevznjJa2v+ihV8uLz+5saduztLelejdwZAeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJDz6IzOdcb5nx3vNhXJAu2yNzToWR8e9SpUjO5AEYQeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgrADSRB2IAnCDiTRMOy2F9p+yvZu2y/Y/nSx/HbbB21vL/5Wtr9dAM2ayPQDw5LWRsTztmdK2mZ7U1G7OyK+2r72ANRlIvOzD0oaLO4ft71H0oJ2NwagXmf0md32JZKWS9paLLrF9g7b623PKlmn3/aA7YFTOtFSswCaN+Gw2z5P0sOSPhMRxyR9XdKlkpZp9Mh/53jrRcS6iOiLiL4pmlZDywCaMaGw256i0aDfGxGPSFJEHIqIkYh4TdI3JK1oX5sAWjWRs/GWdI+kPRFx15jl88c87SOSdtXfHoC6TORs/Psl3Shpp+3txbLbJK22vUxSSNov6ZNt6RBALSZyNv4ZSeP9DvXj9bcDoF34Bh2QBGEHkiDsQBKEHUiCsANJEHYgCcIOJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJR0Tndma/IumlMYvmSDrSsQbOTK/21qt9SfTWrDp7uzgi5o5X6GjY37RzeyAi+rrWQIVe7a1X+5LorVmd6o238UAShB1IotthX9fl/Vfp1d56tS+J3prVkd66+pkdQOd0+8gOoEMIO5BEV8Ju+zrbP7W9z/at3eihjO39tncW01APdLmX9bYP2941Ztls25ts7y1ux51jr0u99cQ03hXTjHf1tev29Ocd/8xue7Kkn0n6gKQDkp6TtDoidne0kRK290vqi4iufwHD9u9KelXStyPiXcWyv5F0NCLuKP6hnBURn+uR3m6X9Gq3p/EuZiuaP3aacUnXS7pJXXztKvq6QR143bpxZF8haV9EvBgRJyU9IGlVF/roeRHxtKSjb1i8StKG4v4Gjf7P0nElvfWEiBiMiOeL+8clnZ5mvKuvXUVfHdGNsC+Q9Isxjw+ot+Z7D0lP2N5mu7/bzYxjXkQMFvdfljSvm82Mo+E03p30hmnGe+a1a2b681Zxgu7NroqId0v6kKRPFW9Xe1KMfgbrpbHTCU3j3SnjTDP+K9187Zqd/rxV3Qj7QUkLxzy+qFjWEyLiYHF7WNKj6r2pqA+dnkG3uD3c5X5+pZem8R5vmnH1wGvXzenPuxH25yQttv1221MlfVTSxi708Sa2ZxQnTmR7hqQPqvemot4oaU1xf42kx7rYy+v0yjTeZdOMq8uvXdenP4+Ijv9JWqnRM/I/l/RX3eihpK9Fkn5c/L3Q7d4k3a/Rt3WnNHpu42ZJvyZps6S9kp6UNLuHevsnSTsl7dBosOZ3qberNPoWfYek7cXfym6/dhV9deR14+uyQBKcoAOSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJP4fcKgKSEIBgPIAAAAASUVORK5CYII=\n"
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Ground truth labels: [0. 0. 1. 0. 0. 0. 0. 0. 0. 0.]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tHq96SIJcNfx",
        "colab_type": "code",
        "colab": {},
        "cellView": "both"
      },
      "source": [
        "#@title Define MNIST model architecture using tf.keras API\n",
        "\n",
        "def simple_mnist_model(input_shape):\n",
        "  \"\"\"Creates a simple (multi-layer perceptron) MNIST model.\"\"\"\n",
        "\n",
        "  model = tf.keras.models.Sequential()\n",
        "  # Flatten to a 1d array (e.g. 28x28 -> 784)\n",
        "  model.add(tf.keras.layers.Flatten(input_shape=input_shape))\n",
        "  # Fully-connected neural layer with 128 neurons, RELU activation\n",
        "  model.add(tf.keras.layers.Dense(128, activation='relu'))\n",
        "  # Fully-connected neural layer returning probability scores for each class\n",
        "  model.add(tf.keras.layers.Dense(10, activation='softmax'))\n",
        "  return model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7Gdxh7qWcPSO",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 374
        },
        "cellView": "form",
        "outputId": "e38d9b12-c27c-4174-824b-8bc232a2c417"
      },
      "source": [
        "#@title Train the Keras model\n",
        "\n",
        "with strategy.scope():\n",
        "  model = simple_mnist_model(input_shape)\n",
        "  tf.print(\"Constructed Keras MNIST model, training...\")\n",
        "\n",
        "  optimizer = tf.keras.optimizers.SGD(learning_rate=0.05)\n",
        "  training_loss = tf.keras.metrics.Mean(\"training_loss\", dtype=tf.float32)\n",
        "  training_accuracy = tf.keras.metrics.CategoricalAccuracy(\n",
        "      \"training_accuracy\", dtype=tf.float32)\n",
        "  test_loss = tf.keras.metrics.Mean(\"test_loss\", dtype=tf.float32)\n",
        "  test_accuracy = tf.keras.metrics.CategoricalAccuracy(\n",
        "      \"test_accuracy\", dtype=tf.float32)\n",
        "\n",
        "  @tf.function\n",
        "  def train_step(iterator):\n",
        "    \"\"\"Training StepFn.\"\"\"\n",
        "\n",
        "    def step_fn(inputs):\n",
        "      \"\"\"Per-Replica StepFn.\"\"\"\n",
        "      images, labels = inputs\n",
        "      with tf.GradientTape() as tape:\n",
        "        logits = model(images, training=True)\n",
        "        loss = tf.keras.losses.categorical_crossentropy(labels, logits)\n",
        "        loss = tf.reduce_mean(loss) / strategy.num_replicas_in_sync\n",
        "      grads = tape.gradient(loss, model.trainable_variables)\n",
        "      optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
        "      training_loss.update_state(loss)\n",
        "      training_accuracy.update_state(labels, logits)\n",
        "\n",
        "    strategy.run(step_fn, args=(next(iterator),))\n",
        "\n",
        "  @tf.function\n",
        "  def test_step(iterator):\n",
        "    \"\"\"Evaluation StepFn.\"\"\"\n",
        "\n",
        "    def step_fn(inputs):\n",
        "      images, labels = inputs\n",
        "      logits = model(images, training=False)\n",
        "      loss = tf.keras.losses.categorical_crossentropy(labels, logits)\n",
        "      loss = tf.reduce_mean(loss) / strategy.num_replicas_in_sync\n",
        "      test_loss.update_state(loss)\n",
        "      test_accuracy.update_state(labels, logits)\n",
        "\n",
        "    strategy.run(step_fn, args=(next(iterator),))\n",
        "\n",
        "  for epoch in range(0, num_epochs):\n",
        "    tf.print(\"Running epoch #%s\" % (epoch + 1))\n",
        "\n",
        "    train_iterator = iter(train_dist_ds)\n",
        "    for step in range(steps_per_epoch):\n",
        "      train_step(train_iterator)\n",
        "    tf.print(\"  Training loss: %f, accuracy: %f\" % (training_loss.result(), training_accuracy.result() * 100))\n",
        "    training_loss.reset_states()\n",
        "    training_accuracy.reset_states()\n",
        "\n",
        "    test_iterator = iter(test_dist_ds)\n",
        "    for step in range(steps_per_eval):\n",
        "      test_step(test_iterator)\n",
        "    tf.print(\"  Test loss    : %f, accuracy: %f\" % (test_loss.result(), test_accuracy.result() * 100))\n",
        "    test_loss.reset_states()\n",
        "    test_accuracy.reset_states()\n",
        "\n",
        "  tf.print(\"Completed training!\")\n",
        "  tf.print(\"\")\n",
        "\n",
        "  # Run a single prediction on the trained model\n",
        "  tf_prediction = model(sample_image_batch, training=False)\n",
        "  tf.print(\"Sample prediction:\")\n",
        "  tf.print(tf_prediction[0] * 100.0, summarize=100)\n",
        "  tf.print(\"\")"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Constructed Keras MNIST model, training...\n",
            "Running epoch #1\n",
            "  Training loss: 0.714714, accuracy: 82.386665\n",
            "  Test loss    : 0.386100, accuracy: 89.630005\n",
            "Running epoch #2\n",
            "  Training loss: 0.361406, accuracy: 90.100006\n",
            "  Test loss    : 0.312490, accuracy: 91.000000\n",
            "Running epoch #3\n",
            "  Training loss: 0.308521, accuracy: 91.355003\n",
            "  Test loss    : 0.278433, accuracy: 92.000000\n",
            "Running epoch #4\n",
            "  Training loss: 0.277716, accuracy: 92.261665\n",
            "  Test loss    : 0.255600, accuracy: 92.650002\n",
            "Running epoch #5\n",
            "  Training loss: 0.255256, accuracy: 92.849998\n",
            "  Test loss    : 0.237870, accuracy: 93.339996\n",
            "Completed training!\n",
            "\n",
            "Sample prediction:\n",
            "[0.942700624 0.0189254563 94.8591919 1.60412228 4.52805943e-06 1.21846104 0.845801711 3.33044341e-06 0.510759115 2.41755624e-05]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DmespEaFcSEL",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 122
        },
        "outputId": "07ba7fb8-bd99-495d-941c-17320014573b"
      },
      "source": [
        "#@title Export the trained model as a SavedModel, with IREE-compatible settings\n",
        "\n",
        "# Since the model was written in sequential style, explicitly wrap in a module.\n",
        "saved_model_dir = \"/tmp/mnist.sm\"\n",
        "inference_module = tf.Module()\n",
        "inference_module.model = model\n",
        "# Hack: Convert to static shape. Won't be necessary once dynamic shapes are in.\n",
        "dynamic_input_shape = list(model.inputs[0].shape)\n",
        "dynamic_input_shape[0] = 1  # Make fixed (batch=1)\n",
        "# Produce a concrete function.\n",
        "inference_module.predict = tf.function(\n",
        "    input_signature=[\n",
        "        tf.TensorSpec(dynamic_input_shape, model.inputs[0].dtype)])(\n",
        "            lambda x: model.call(x, training=False))\n",
        "save_options = tf.saved_model.SaveOptions(save_debug_info=True)\n",
        "tf.print(\"Exporting SavedModel to %s\" % saved_model_dir)\n",
        "tf.saved_model.save(inference_module, saved_model_dir, options=save_options)"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Exporting SavedModel to /tmp/mnist.sm\n",
            "WARNING:tensorflow:From /usr/local/lib/python3.7/dist-packages/tensorflow_core/python/ops/resource_variable_ops.py:1788: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "If using Keras pass *_constraint arguments to layers.\n",
            "INFO:tensorflow:Assets written to: /tmp/mnist.sm/assets\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nZdVUd_dgTtc",
        "colab_type": "text"
      },
      "source": [
        "# Compile and Execute MNIST Model using IREE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rqwIx4j4gS1a",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 717
        },
        "outputId": "2a9b7627-8836-427d-af51-5e001c6a9226"
      },
      "source": [
        "#@title Load the SavedModel into IREE's compiler as MLIR xla_hlo\n",
        "\n",
        "compiler_module = ireec.tf_load_saved_model(\n",
        "    saved_model_dir, exported_names=[\"predict\"])\n",
        "tf.print(\"Imported MLIR:\\n\", compiler_module.to_asm(large_element_limit=100))\n",
        "\n",
        "# Write to a file for use outside of this notebook.\n",
        "mnist_mlir_path = os.path.join(SAVE_PATH, \"mnist.mlir\")\n",
        "with open(mnist_mlir_path, \"wt\") as output_file:\n",
        "  output_file.write(compiler_module.to_asm())\n",
        "print(\"Wrote MLIR to path '%s'\" % mnist_mlir_path)"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Imported MLIR:\n",
            " \n",
            "\n",
            "module {\n",
            "  flow.variable @\"__iree_flow___sm_node15__model.layer-2.kernel\" opaque<\"\", \"0xDEADBEEF\"> : tensor<784x128xf32>\n",
            "  flow.variable @\"__iree_flow___sm_node16__model.layer-2.bias\" opaque<\"\", \"0xDEADBEEF\"> : tensor<128xf32>\n",
            "  flow.variable @\"__iree_flow___sm_node21__model.layer-3.kernel\" opaque<\"\", \"0xDEADBEEF\"> : tensor<128x10xf32>\n",
            "  flow.variable @\"__iree_flow___sm_node22__model.layer-3.bias\" dense<[-0.124819458, 0.10169369, -0.0273698159, -0.0204691291, 0.0623821244, 0.213322133, -0.036071334, 0.112993836, -0.235609874, -0.0460525416]> : tensor<10xf32>\n",
            "  func @predict(%arg0: tensor<1x28x28x1xf32>) -> tensor<1x10xf32> attributes {iree.module.export, iree.reflection = {abi = \"sip\", abiv = 1 : i32, sip = \"I8!S5!k0_0R3!_0\"}, tf._input_shapes = [\"tfshape$dim { size: 1 } dim { size: 28 } dim { size: 28 } dim { size: 1 }\", \"tfshape$unknown_rank: true\", \"tfshape$unknown_rank: true\", \"tfshape$unknown_rank: true\", \"tfshape$unknown_rank: true\"], tf.signature.is_stateful} {\n",
            "    %0 = flow.variable.load @\"__iree_flow___sm_node22__model.layer-3.bias\" : tensor<10xf32>\n",
            "    %1 = flow.variable.load @\"__iree_flow___sm_node21__model.layer-3.kernel\" : tensor<128x10xf32>\n",
            "    %2 = flow.variable.load @\"__iree_flow___sm_node16__model.layer-2.bias\" : tensor<128xf32>\n",
            "    %3 = flow.variable.load @\"__iree_flow___sm_node15__model.layer-2.kernel\" : tensor<784x128xf32>\n",
            "    %4 = xla_hlo.constant opaque<\"\", \"0xDEADBEEF\"> : tensor<1x128xf32>\n",
            "    %5 = xla_hlo.constant dense<0xFF800000> : tensor<f32>\n",
            "    %6 = xla_hlo.constant dense<0.000000e+00> : tensor<f32>\n",
            "    %7 = \"xla_hlo.reshape\"(%arg0) : (tensor<1x28x28x1xf32>) -> tensor<1x784xf32>\n",
            "    %8 = \"xla_hlo.dot\"(%7, %3) : (tensor<1x784xf32>, tensor<784x128xf32>) -> tensor<1x128xf32>\n",
            "    %9 = \"xla_hlo.add\"(%8, %2) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<1x128xf32>, tensor<128xf32>) -> tensor<1x128xf32>\n",
            "    %10 = xla_hlo.maximum %9, %4 : tensor<1x128xf32>\n",
            "    %11 = \"xla_hlo.dot\"(%10, %1) : (tensor<1x128xf32>, tensor<128x10xf32>) -> tensor<1x10xf32>\n",
            "    %12 = \"xla_hlo.add\"(%11, %0) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<1x10xf32>, tensor<10xf32>) -> tensor<1x10xf32>\n",
            "    %13 = \"xla_hlo.reduce\"(%12, %5) ( {\n",
            "    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):\t// no predecessors\n",
            "      %18 = xla_hlo.maximum %arg1, %arg2 : tensor<f32>\n",
            "      \"xla_hlo.return\"(%18) : (tensor<f32>) -> ()\n",
            "    }) {dimensions = dense<1> : tensor<1xi64>} : (tensor<1x10xf32>, tensor<f32>) -> tensor<1xf32>\n",
            "    %14 = \"xla_hlo.subtract\"(%12, %13) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<1x10xf32>, tensor<1xf32>) -> tensor<1x10xf32>\n",
            "    %15 = \"xla_hlo.exponential\"(%14) : (tensor<1x10xf32>) -> tensor<1x10xf32>\n",
            "    %16 = \"xla_hlo.reduce\"(%15, %6) ( {\n",
            "    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):\t// no predecessors\n",
            "      %18 = xla_hlo.add %arg1, %arg2 : tensor<f32>\n",
            "      \"xla_hlo.return\"(%18) : (tensor<f32>) -> ()\n",
            "    }) {dimensions = dense<1> : tensor<1xi64>} : (tensor<1x10xf32>, tensor<f32>) -> tensor<1xf32>\n",
            "    %17 = \"xla_hlo.divide\"(%15, %16) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<1x10xf32>, tensor<1xf32>) -> tensor<1x10xf32>\n",
            "    return %17 : tensor<1x10xf32>\n",
            "  }\n",
            "}\n",
            "\n",
            "Wrote MLIR to path '/usr/local/google/home/scotttodd/saved_models/mnist.mlir'\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IDHI7h3khJr9",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "f0676fc3-d7e7-497e-f28d-480e33c37fc9"
      },
      "source": [
        "#@title Compile the xla_hlo MLIR and prepare a context to execute it\n",
        "\n",
        "# Compile the MLIR module into a VM module for execution\n",
        "flatbuffer_blob = compiler_module.compile(target_backends=[backend_name])\n",
        "vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob)\n",
        "\n",
        "# Register the module with a runtime context\n",
        "config = ireert.Config(driver_name)\n",
        "ctx = ireert.SystemContext(config=config)\n",
        "ctx.add_module(vm_module)"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Created IREE driver vulkan: <pyiree.rt.binding.HalDriver object at 0x7f6e4471bc70>\n",
            "SystemContext driver=<pyiree.rt.binding.HalDriver object at 0x7f6e4471bc70>\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SKflpnLtkLYE",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 102
        },
        "outputId": "48996e9a-f70f-4f45-b53c-fd9fab80d78d"
      },
      "source": [
        "#@title Execute the compiled module and compare the results with TensorFlow\n",
        "\n",
        "# Invoke the 'predict' function with a single image as an argument\n",
        "iree_prediction = ctx.modules.module.predict(sample_image_batch)\n",
        "\n",
        "tf.print(\"IREE prediction ('%s' backend, '%s' driver):\" % (backend_name, driver_name))\n",
        "tf.print(tf.convert_to_tensor(iree_prediction[0]) * 100.0, summarize=100)\n",
        "tf.print(\"\")\n",
        "tf.print(\"TensorFlow prediction:\")\n",
        "tf.print(tf_prediction[0] * 100.0, summarize=100)"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "IREE prediction ('vulkan-spirv' backend, 'vulkan' driver):\n",
            "[0.942700088 0.018925447 94.8592 1.60412395 4.52806762e-06 1.2184602 0.845801711 3.33043749e-06 0.510759652 2.41755624e-05]\n",
            "\n",
            "TensorFlow prediction:\n",
            "[0.942700624 0.0189254563 94.8591919 1.60412228 4.52805943e-06 1.21846104 0.845801711 3.33044341e-06 0.510759115 2.41755624e-05]\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}