TPUs_in_Colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "TPUs in Colab",
      "provenance": [],
      "collapsed_sections": [
        "clSFHJkFNylD"
      ],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/hjmus/mock/blob/master/TPUs_in_Colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ovFDeMgtjqW4"
      },
      "source": [
        "# TPUs in Colab&nbsp; <a href=\"https://cloud.google.com/tpu/\"><img valign=\"middle\" src=\"https://raw.githubusercontent.com/GoogleCloudPlatform/tensorflow-without-a-phd/master/tensorflow-rl-pong/images/tpu-hexagon.png\" width=\"50\"></a>\n",
        "In this example, we'll work through training a model to classify images of\n",
        "flowers on Google's lightning-fast Cloud TPUs. Our model will take as input a photo of a flower and return whether it is a daisy, dandelion, rose, sunflower, or tulip.\n",
        "\n",
        "We use the Keras framework, new to TPUs in TF 2.1.0. Adapted from [this notebook](https://colab.research.google.com/github/GoogleCloudPlatform/training-data-analyst/blob/master/courses/fast-and-lean-data-science/07_Keras_Flowers_TPU_xception_fine_tuned_best.ipynb) by [Martin Gorner](https://twitter.com/martin_gorner)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "clSFHJkFNylD"
      },
      "source": [
        "#### License"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hleIN5-pcr0N"
      },
      "source": [
        "Copyright 2019-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",
        "    http://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.\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "This is not an official Google product but sample code provided for an educational purpose.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_pQCOmISAQBu"
      },
      "source": [
        "## Enabling and testing the TPU\n",
        "\n",
        "First, you'll need to enable TPUs for the notebook:\n",
        "\n",
        "- Navigate to Edit→Notebook Settings\n",
        "- select TPU from the Hardware Accelerator drop-down\n",
        "\n",
        "Next, we'll check that we can connect to the TPU:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FpvUOuC3j27n",
        "outputId": "2e16870d-5caa-477b-e840-fa78a5609101"
      },
      "source": [
        "%tensorflow_version 2.x\n",
        "import tensorflow as tf\n",
        "print(\"Tensorflow version \" + tf.__version__)\n",
        "\n",
        "try:\n",
        "  tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection\n",
        "  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])\n",
        "except ValueError:\n",
        "  raise BaseException('ERROR: Not connected to a TPU runtime; please see the previous cell in this notebook for instructions!')\n",
        "\n",
        "tf.config.experimental_connect_to_cluster(tpu)\n",
        "tf.tpu.experimental.initialize_tpu_system(tpu)\n",
        "tpu_strategy = tf.distribute.experimental.TPUStrategy(tpu)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Tensorflow version 2.5.0\n",
            "Running on TPU  ['10.34.6.242:8470']\n",
            "WARNING:tensorflow:TPU system grpc://10.34.6.242:8470 has already been initialized. Reinitializing the TPU can cause previously created variables on TPU to be lost.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:TPU system grpc://10.34.6.242:8470 has already been initialized. Reinitializing the TPU can cause previously created variables on TPU to be lost.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Initializing the TPU system: grpc://10.34.6.242:8470\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Initializing the TPU system: grpc://10.34.6.242:8470\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Clearing out eager caches\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Clearing out eager caches\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Finished initializing TPU system.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Finished initializing TPU system.\n",
            "WARNING:absl:`tf.distribute.experimental.TPUStrategy` is deprecated, please use  the non experimental symbol `tf.distribute.TPUStrategy` instead.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Found TPU system:\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Found TPU system:\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores: 8\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores: 8\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Workers: 1\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Workers: 1\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OH1M0Qi9kBUO",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 366
        },
        "outputId": "557be515-0d70-4d15-f156-216a2c2c977f"
      },
      "source": [
        "from tensorflow.compiler.xla.experimental.xla_sharding import xla_sharding\n",
        "import numpy as np\n",
        "\n",
        "with tpu_strategy.scope():\n",
        "  optimizer = tf.keras.optimizers.Adam()\n",
        "  var0 = tf.Variable([4.0, 5.0])\n",
        "  var0_split = tf.Variable([14.0, 15.0])  # placeholder\n",
        "  a = tf.constant(4.0)\n",
        "\n",
        "dev_mesh = np.array([0, 1])\n",
        "tensor_split_dims_mapping = [0]\n",
        "\n",
        "@tf.function(jit_compile=True)\n",
        "def test1(x):\n",
        "    loss = x * var0\n",
        "    print('loss', loss)\n",
        "\n",
        "    return tf.reduce_sum(loss)\n",
        "\n",
        "# print(test1.experimental_get_compiler_ir(a)(stage='hlo'))\n",
        "print(test1.experimental_get_compiler_ir(a)(stage='hlo', device_name='/job:worker/replica:0/task:0/device:TPU:1'))\n"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "loss Tensor(\"mul:0\", shape=(2,), dtype=float32)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-10-55d9c9579ae1>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     19\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     20\u001b[0m \u001b[0;31m# print(test1.experimental_get_compiler_ir(a)(stage='hlo'))\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 21\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtest1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexperimental_get_compiler_ir\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstage\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'hlo'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice_name\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'/job:worker/replica:0/task:0/device:TPU:1'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36mcompiler_ir_generator\u001b[0;34m(stage, device_name)\u001b[0m\n\u001b[1;32m   1102\u001b[0m           \u001b[0mstage\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mstage\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1103\u001b[0m           \u001b[0mfunction_name\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfn_name\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1104\u001b[0;31m           args=list(filtered_flat_args) + concrete_fn.captured_inputs)\n\u001b[0m\u001b[1;32m   1105\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0mstage\u001b[0m \u001b[0;32min\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m\"hlo_serialized\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"optimized_hlo_serialized\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1106\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mres_bytes\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/context.py\u001b[0m in \u001b[0;36mget_compiler_ir\u001b[0;34m(self, device_name, function_name, args, stage)\u001b[0m\n\u001b[1;32m   1527\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mget_compiler_ir\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfunction_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstage\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"hlo\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1528\u001b[0m     return pywrap_tfe.TF_GetCompilerIr(self._context_handle, function_name,\n\u001b[0;32m-> 1529\u001b[0;31m                                        stage, device_name, args)\n\u001b[0m\u001b[1;32m   1530\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1531\u001b[0m   @deprecated(\n",
            "\u001b[0;31mValueError\u001b[0m: No matching device found for '/job:worker/replica:0/task:0/device:TPU:1'"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CuxDFfGbxmZP",
        "outputId": "0ae8d62f-aac0-464d-d581-59cc2a59a488",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 401
        }
      },
      "source": [
        "@tf.function(jit_compile=True)\n",
        "def test_xla_sharding1(x):\n",
        "    var0_split.assign(xla_sharding.mesh_split(var0, dev_mesh, tensor_split_dims_mapping, use_sharding_op=True))\n",
        "    loss = x * var0_split\n",
        "    print('loss', loss)\n",
        "\n",
        "    return tf.reduce_sum(loss)\n",
        "\n",
        "\n",
        "@tf.function(jit_compile=True)\n",
        "def test_xla_sharding2(x):\n",
        "    with tf.GradientTape() as tape:\n",
        "        var0_split.assign(xla_sharding.mesh_split(var0, dev_mesh, tensor_split_dims_mapping, use_sharding_op=True))\n",
        "        loss = x * var0_split\n",
        "        print('loss', loss)\n",
        "    grads = tape.gradient(loss, [var0_split])\n",
        "    print('grads', grads)\n",
        "    var1 = var0_split - grads\n",
        "\n",
        "    return tf.reduce_sum(var1)\n",
        "\n",
        "@tf.function(jit_compile=True)\n",
        "def test_xla_sharding3(x):\n",
        "    with tf.GradientTape() as tape:\n",
        "        var0_split.assign(xla_sharding.mesh_split(var0, dev_mesh, tensor_split_dims_mapping, use_sharding_op=True))\n",
        "        print('var0_split', var0_split)\n",
        "        loss = x * var0_split\n",
        "        print('loss', loss)\n",
        "    grads = tape.gradient(loss, [var0_split])\n",
        "    print('grads', grads)\n",
        "    optimizer.apply_gradients(zip(grads, [var0_split]))\n",
        "    print('applied gradients')\n",
        "\n",
        "    return tf.reduce_sum(var0_split)\n",
        "\n",
        "test_xla_sharding1(a)\n",
        "\n",
        "print(test_xla_sharding1.experimental_get_compiler_ir(a)(stage='hlo', device_name='/job:worker/replica:0/task:0/device:TPU:1'))\n",
        "print(test_xla_sharding2.experimental_get_compiler_ir(a)(stage='hlo', device_name='/job:worker/replica:0/task:0/device:TPU:1'))\n",
        "print(test_xla_sharding3.experimental_get_compiler_ir(a)(stage='hlo', device_name='/job:worker/replica:0/task:0/device:TPU:1'))\n",
        "\n"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "loss Tensor(\"mul:0\", shape=(2,), dtype=float32)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-11-a8d7f341f7e2>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     36\u001b[0m \u001b[0mtest_xla_sharding1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     37\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 38\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtest_xla_sharding1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexperimental_get_compiler_ir\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstage\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'hlo'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice_name\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'/job:worker/replica:0/task:0/device:TPU:1'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     39\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtest_xla_sharding2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexperimental_get_compiler_ir\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstage\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'hlo'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice_name\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'/job:worker/replica:0/task:0/device:TPU:1'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtest_xla_sharding3\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexperimental_get_compiler_ir\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstage\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'hlo'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice_name\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'/job:worker/replica:0/task:0/device:TPU:1'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36mcompiler_ir_generator\u001b[0;34m(stage, device_name)\u001b[0m\n\u001b[1;32m   1102\u001b[0m           \u001b[0mstage\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mstage\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1103\u001b[0m           \u001b[0mfunction_name\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfn_name\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1104\u001b[0;31m           args=list(filtered_flat_args) + concrete_fn.captured_inputs)\n\u001b[0m\u001b[1;32m   1105\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0mstage\u001b[0m \u001b[0;32min\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m\"hlo_serialized\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"optimized_hlo_serialized\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1106\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mres_bytes\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/context.py\u001b[0m in \u001b[0;36mget_compiler_ir\u001b[0;34m(self, device_name, function_name, args, stage)\u001b[0m\n\u001b[1;32m   1527\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mget_compiler_ir\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfunction_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstage\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"hlo\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1528\u001b[0m     return pywrap_tfe.TF_GetCompilerIr(self._context_handle, function_name,\n\u001b[0;32m-> 1529\u001b[0;31m                                        stage, device_name, args)\n\u001b[0m\u001b[1;32m   1530\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1531\u001b[0m   @deprecated(\n",
            "\u001b[0;31mValueError\u001b[0m: No matching device found for '/job:worker/replica:0/task:0/device:TPU:1'"
          ]
        }
      ]
    }
  ]
}