<commit_before>{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title 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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "# Use TPUs\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/tpu\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ys81cOhXOWUP"
      },
      "source": [
        "Experimental support for Cloud TPUs is currently available for Keras\n",
        "and Google Colab. Before you run this Colab notebooks, ensure that\n",
        "your hardware accelerator is a TPU by checking your notebook settings:\n",
        "Runtime > Change runtime type > Hardware accelerator > TPU."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ek5Hop74NVKm"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cw0WRaChRxTL"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "import os\n",
        "import tensorflow_datasets as tfds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yDWaRxSpwBN1"
      },
      "source": [
        "## TPU Initialization\n",
        "TPUs are usually on Cloud TPU workers which are different from the local process running the user python program. Thus some initialization work needs to be done to connect to the remote cluster and initialize TPUs. Note that the `tpu` argument to `TPUClusterResolver` is a special address just for Colab. In the case that you are running on Google Compute Engine (GCE), you should instead pass in the name of your CloudTPU."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dCqWMqvtwOLs"
      },
      "source": [
        "Note: The TPU initialization code has to be at the beginning of your program."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dKPqF8d1wJCV"
      },
      "outputs": [],
      "source": [
        "resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])\n",
        "tf.config.experimental_connect_to_cluster(resolver)\n",
        "# This is the TPU initialization code that has to be at the beginning.\n",
        "tf.tpu.experimental.initialize_tpu_system(resolver)\n",
        "print(\"All devices: \", tf.config.list_logical_devices('TPU'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mv7kehTZ1Lq_"
      },
      "source": [
        "## Manual device placement\n",
        "After the TPU is initialized, you can use manual device placement to place the computation on a single TPU device.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRZ4kMoxBNND"
      },
      "outputs": [],
      "source": [
        "a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])\n",
        "b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])\n",
        "with tf.device('/TPU:0'):\n",
        "  c = tf.matmul(a, b)\n",
        "print(\"c device: \", c.device)\n",
        "print(c)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_NJm-kgFO0cC"
      },
      "source": [
        "## Distribution strategies\n",
        "Most times users want to run the model on multiple TPUs in a data parallel way. A distribution strategy is an abstraction that can be used to drive models on CPU, GPUs or TPUs. Simply swap out the distribution strategy and the model will run on the given device. See the [distribution strategy guide](./distributed_training.ipynb) for more information."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DcDPMZs-9uLJ"
      },
      "source": [
        "First, creates the `TPUStrategy` object."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7SO23K8oRpjI"
      },
      "outputs": [],
      "source": [
        "strategy = tf.distribute.TPUStrategy(resolver)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JlaAmswWPsU6"
      },
      "source": [
        "To replicate a computation so it can run in all TPU cores, you can simply pass it to `strategy.run` API. Below is an example that all the cores will obtain the same inputs `(a, b)`, and do the matmul on each core independently. The outputs will be the values from all the replicas."
      ]
    },
    {<commit_msg>Translate tpu.ipynb via GitLocalize<commit_after>{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title 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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "# Use TPUs\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://www.tensorflow.org/guide/tpu\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\">View on TensorFlow.org</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Run in Google Colab</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">View source on GitHub</a></td>\n",
        "  <td><a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\">Download notebook</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ys81cOhXOWUP"
      },
      "source": [
        "Experimental support for Cloud TPUs is currently available for Keras\n",
        "and Google Colab. Before you run this Colab notebooks, ensure that\n",
        "your hardware accelerator is a TPU by checking your notebook settings:\n",
        "Runtime > Change runtime type > Hardware accelerator > TPU."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ek5Hop74NVKm"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cw0WRaChRxTL"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "import os\n",
        "import tensorflow_datasets as tfds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yDWaRxSpwBN1"
      },
      "source": [
        "## TPU Initialization\n",
        "TPUs are usually on Cloud TPU workers which are different from the local process running the user python program. Thus some initialization work needs to be done to connect to the remote cluster and initialize TPUs. Note that the `tpu` argument to `TPUClusterResolver` is a special address just for Colab. In the case that you are running on Google Compute Engine (GCE), you should instead pass in the name of your CloudTPU."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dCqWMqvtwOLs"
      },
      "source": [
        "Note: The TPU initialization code has to be at the beginning of your program."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dKPqF8d1wJCV"
      },
      "outputs": [],
      "source": [
        "resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])\n",
        "tf.config.experimental_connect_to_cluster(resolver)\n",
        "# This is the TPU initialization code that has to be at the beginning.\n",
        "tf.tpu.experimental.initialize_tpu_system(resolver)\n",
        "print(\"All devices: \", tf.config.list_logical_devices('TPU'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mv7kehTZ1Lq_"
      },
      "source": [
        "## Manual device placement\n",
        "After the TPU is initialized, you can use manual device placement to place the computation on a single TPU device.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRZ4kMoxBNND"
      },
      "outputs": [],
      "source": [
        "a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])\n",
        "b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])\n",
        "with tf.device('/TPU:0'):\n",
        "  c = tf.matmul(a, b)\n",
        "print(\"c device: \", c.device)\n",
        "print(c)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_NJm-kgFO0cC"
      },
      "source": [
        "## Distribution strategies\n",
        "Most times users want to run the model on multiple TPUs in a data parallel way. A distribution strategy is an abstraction that can be used to drive models on CPU, GPUs or TPUs. Simply swap out the distribution strategy and the model will run on the given device. See the [distribution strategy guide](./distributed_training.ipynb) for more information."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DcDPMZs-9uLJ"
      },
      "source": [
        "First, creates the `TPUStrategy` object."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7SO23K8oRpjI"
      },
      "outputs": [],
      "source": [
        "strategy = tf.distribute.experimental.TPUStrategy(resolver)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JlaAmswWPsU6"
      },
      "source": [
        "To replicate a computation so it can run in all TPU cores, you can simply pass it to `strategy.run` API. Below is an example that all the cores will obtain the same inputs `(a, b)`, and do the matmul on each core independently. The outputs will be the values from all the replicas."
      ]
    },
    {
{'fraction_non_alphanumeric': 0.15794580876315592, 'fraction_numerical': 0.010002239307307606, 'mean_word_length': 2.1421200750469045, 'pattern_counts': {'":': 144, '<': 47, '<?xml version=': 0, '>': 53, 'https://': 18, 'lorem ipsum': 0, 'www.': 12, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4297649', 'n_tokens_mistral': 4459, 'n_tokens_neox': 4353, 'n_words': 1117}