TF-Coder_Colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "TF-Coder Colab.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "5u_aruynnBDo",
        "EKWOuU6Ilr-3",
        "pHYlFkj6l5af",
        "Q6uRr4x9WHRC",
        "NYWv_mSbmfOO"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5u_aruynnBDo"
      },
      "source": [
        "##### Copyright 2020 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xKqvIgn4nFP7"
      },
      "source": [
        "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": "jSKzsPs5nMtg"
      },
      "source": [
        "# TensorFlow Coder (TF-Coder): A program synthesis tool for TensorFlow expressions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Myqm791nURE"
      },
      "source": [
        "**TensorFlow Coder** is a tool that helps you manipulate tensors with TensorFlow! If you provide an example of a tensor manipulation, TF-Coder will search for TensorFlow code that matches the example.\n",
        "\n",
        "Follow this [**tutorial**](https://github.com/google-research/tensorflow-coder/blob/master/Tutorial.md) to get familiar with TF-Coder.\n",
        "\n",
        "Make sure to connect to a runtime (click \"Connect\" in the top right corner)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Iq76A8wwfMpC"
      },
      "source": [
        "## Step 0: Data collection request"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RyKI7_BpfSYl"
      },
      "source": [
        "**Note from the TF-Coder team at Google:**\n",
        "\n",
        "We are excited to bring you TF-Coder, which we hope will accelerate your TensorFlow development.\n",
        "\n",
        "We have one quick request first: we would like to log usage data for TF-Coder, so that we can identify scenarios where TF-Coder can be improved. This usage data will help us improve TF-Coder for everyone. We also believe that the usage data will be a valuable resource to the broader program synthesis research community. Please read the text below and then use the following cell to let us know whether we may log or release your usage data. Either way, you may still use the TF-Coder tool.\n",
        "\n",
        "---\n",
        "\n",
        "Collecting TF-Coder usage data will help Google improve the TF-Coder tool, and TensorFlow services more generally. This usage data includes (i) the problems you create, (ii) the settings for the TF-Coder tool, (iii) the TF-Coder tool's results for those problems, (iv) metadata relating to your session, problem and device you are using to use the TF-Coder tool, and (v) your location (determined by your IP address). The usage data does not include any other personally identifiable information. Please do not upload or provide any personal or confidential information to the TF-Coder tool.\n",
        "\n",
        "In addition to Google’s internal use of your usage data, Google would also like to release some of such data in a public dataset to facilitate related research and to promote reproducible research publications. If your usage data is released it will be done in an open source fashion, meaning anyone with access to the data may use it for their purposes.\n",
        "\n",
        "To opt-out of Google collecting your usage data entirely, uncheck the first box in the cell below. To opt out of your usage data being released as part of a public dataset, uncheck the second box. For the avoidance of doubt, if you only uncheck the second box, you are consenting to Google’s internal use of your usage data consistent with this disclosure. Regardless of your choice about sharing your usage data, you may still access and use the TF-Coder tool."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Kc0H6a_dKLnb",
        "cellView": "form"
      },
      "source": [
        "#@title Run this cell after making your choices.\n",
        "\n",
        "allow_data_collection = True  #@param {type: \"boolean\"}\n",
        "include_in_dataset = True  #@param {type: \"boolean\"}\n",
        "\n",
        "if allow_data_collection:\n",
        "  if include_in_dataset:\n",
        "    print('Usage data may be collected and released in a public dataset.')\n",
        "  else:\n",
        "    print('Usage data may be collected but will not be publicly released.')\n",
        "else:\n",
        "  print('Usage data will not be collected.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UeQBWIOUgZSH"
      },
      "source": [
        "## Step 1: Installs and imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4K3bU0aSMEwK",
        "cellView": "form"
      },
      "source": [
        "#@title Run this cell to install and import TF-Coder.\n",
        "\n",
        "ready = True\n",
        "try:\n",
        "  _ = (allow_data_collection, include_in_dataset)\n",
        "except NameError as e:\n",
        "  print('Please run the cell in Step 0 first.')\n",
        "  ready = False\n",
        "\n",
        "if ready:\n",
        "  # Import TensorFlow and NumPy in case the user wants to create the example\n",
        "  # programmatically.\n",
        "  import tensorflow as tf\n",
        "  import numpy as np\n",
        "  \n",
        "  !pip install tensorflow-coder\n",
        "  from tf_coder.value_search import colab_interface\n",
        "  from tf_coder.value_search import value_search_settings as settings_module\n",
        "\n",
        "  if allow_data_collection:\n",
        "    !pip install tensorflow-coder-colab-logging\n",
        "    from tf_coder_colab_logging import colab_logging\n",
        "\n",
        "  from google.colab import output\n",
        "  output.clear()\n",
        "\n",
        "  print('Imports successful. Loading models...')\n",
        "  colab_interface.warm_up()\n",
        "  print('Done. TF-Coder is now ready to use!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o_EbjwYGiU1Y"
      },
      "source": [
        "## Step 2: Describe the problem with an example"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GdTDDuFMiXlp"
      },
      "source": [
        "Provide an **input-output example**:\n",
        "\n",
        "* `inputs` is a dictionary containing one or more input tensors with variable names.\n",
        "* `output` is the corresponding output tensor.\n",
        "\n",
        "Tensors can be provided as lists (possibly multidimensional) or `tf.Tensor` objects.\n",
        "\n",
        "You may also specify relevant **scalar constants**. TF-Coder also uses heuristics to guess a few useful constants.\n",
        "\n",
        "Finally, it often helps to provide an **English description** of the desired tensor manipulation. This description can help the tool decide which TensorFlow operations to prioritize.\n",
        "\n",
        "_Note: Please do not include confidential or personal information._"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lWNAIUNgg3RJ"
      },
      "source": [
        "# Edit this cell! Follow the format of the example below.\n",
        "\n",
        "# A dict mapping input variable names to input tensors.\n",
        "inputs = {\n",
        "    'rows': [10, 20, 30],\n",
        "    'cols': [1, 2, 3, 4],\n",
        "}\n",
        "\n",
        "# The corresponding output tensor.\n",
        "output = [[11, 12, 13, 14],\n",
        "          [21, 22, 23, 24],\n",
        "          [31, 32, 33, 34]]\n",
        "\n",
        "# A list of relevant scalar constants, if any.\n",
        "constants = []\n",
        "\n",
        "# An English description of the tensor manipulation.\n",
        "description = 'add two vectors with broadcasting to get a matrix'"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hZK8-vAhiz6q"
      },
      "source": [
        "## Step 3: Run the TF-Coder tool"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Cn8Ofoo6i1W0",
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f87d2359-84f4-4691-d272-96e9ca50737d"
      },
      "source": [
        "#@title Run this cell to invoke TF-Coder on the problem from Step 2.\n",
        "\n",
        "ready = True\n",
        "try:\n",
        "  _ = colab_interface\n",
        "except NameError:\n",
        "  print('Run the cell in Step 1 first.')\n",
        "  ready = False\n",
        "try:\n",
        "  _ = (inputs, output, constants, description)\n",
        "except NameError:\n",
        "  print('Define the problem by running the cell in Step 2 first.')\n",
        "  ready = False\n",
        "\n",
        "#@markdown &nbsp;\n",
        "#@markdown #### **Settings for TF-Coder**\n",
        "#@markdown How long to search for a solution, in seconds.\n",
        "time_limit = 60  #@param {type: \"integer\"}\n",
        "#@markdown How many solutions to find before stopping. If more than 1, the entire search will slow down.\n",
        "number_of_solutions = 1  #@param{type: \"integer\"}\n",
        "#@markdown Whether solutions must use all inputs, at least one input, or no such requirement.\n",
        "solution_requirement = \"all inputs\" #@param [\"all inputs\", \"one input\", \"no restriction\"]\n",
        "\n",
        "settings = settings_module.from_dict({\n",
        "    'timeout': time_limit,\n",
        "    'only_minimal_solutions': False,\n",
        "    'max_solutions': number_of_solutions,\n",
        "    'require_all_inputs_used': solution_requirement == 'all inputs',\n",
        "    'require_one_input_used': solution_requirement == 'one input',\n",
        "})\n",
        "\n",
        "if ready:\n",
        "  if allow_data_collection:\n",
        "    problem_id = colab_logging.get_uuid()\n",
        "    colab_logging.log_problem(inputs, output, constants, description, settings,\n",
        "                              include_in_dataset=include_in_dataset,\n",
        "                              problem_id=problem_id)\n",
        "\n",
        "  # Results will be printed to the cell's output.\n",
        "  results = colab_interface.run_value_search_from_colab(\n",
        "      inputs, output, constants, description, settings)\n",
        "\n",
        "  if allow_data_collection:\n",
        "    colab_logging.log_result(results,\n",
        "                             include_in_dataset=include_in_dataset,\n",
        "                             problem_id=problem_id)"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Input 'rows':\n",
            "tf.Tensor([10 20 30], shape=(3,), dtype=int32)\n",
            "\n",
            "Input 'cols':\n",
            "tf.Tensor([1 2 3 4], shape=(4,), dtype=int32)\n",
            "\n",
            "Output:\n",
            "tf.Tensor(\n",
            "[[11 12 13 14]\n",
            " [21 22 23 24]\n",
            " [31 32 33 34]], shape=(3, 4), dtype=int32)\n",
            "\n",
            "Constants: [0, 1, -1, True, False, 3, 4]\n",
            "\n",
            "Description: add two vectors with broadcasting to get a matrix\n",
            "\n",
            "Searching...\n",
            "\n",
            "Found solution: tf.add(cols, tf.expand_dims(rows, 1))\n",
            "\n",
            "Solution was found in 0.3 seconds:\n",
            "tf.add(cols, tf.expand_dims(rows, 1))\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ju9GylK-m-r6"
      },
      "source": [
        "# &nbsp;\n",
        "---\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EKWOuU6Ilr-3"
      },
      "source": [
        "## Usage Tips"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iTkbOQf5v23k"
      },
      "source": [
        "#### General"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EqaFf533vxdJ"
      },
      "source": [
        "* If TF-Coder finds a solution, it is _guaranteed_ that the solution produces\n",
        "  the example output when run on the example inputs. However, it is _not\n",
        "  guaranteed_ that the solution generalizes in the way you intend! Please\n",
        "  carefully review solutions produced by TF-Coder before using them in your real\n",
        "  project.\n",
        "\n",
        "* TF-Coder will often produce a solution that uses hardcoded constants for\n",
        "  shapes or lengths, e.g., `tf.reshape(to_flatten, (6,))` in order to flatten an\n",
        "  input tensor with shape `(2, 3)`. You may need to manually change these\n",
        "  constants to improve the generality of the solution, e.g., replacing `6` with\n",
        "  `-1` in this case. Use the shape attribute to obtain dimension lengths of\n",
        "  input tensors, e.g., `to_flatten.shape[0]` would be `2`.\n",
        "\n",
        "* If you want to play with TensorFlow in Colab (e.g., to understand how a\n",
        "  TF-Coder solution works or to test your own solution):\n",
        "  * The TF-Coder Colab already imports TensorFlow 2 and Numpy, for your\n",
        "    convenience.\n",
        "  * Use `tf.constant` to create a tensor from the list format:\n",
        "    ```\n",
        "    >>> tf.constant([[13, 22], [17, 5]])\n",
        "    <tf.Tensor: id=1, shape=(2, 2), dtype=int32, numpy=\n",
        "    array([[13, 22],\n",
        "           [17,  5]], dtype=int32)>\n",
        "\n",
        "    >>> tf.constant(12.3)\n",
        "    <tf.Tensor: id=2, shape=(), dtype=float32, numpy=12.3>\n",
        "    ```\n",
        "  * A Colab notebook can only have one cell running at a time. If you want to\n",
        "    experiment with TensorFlow code while TF-Coder is running, consider doing so\n",
        "    in a separate Python shell.\n",
        "\n",
        "* TF-Coder's running time is exponential in the complexity of the solution.\n",
        "  _Simplifying the problem_, or _breaking it down into multiple steps_, can help\n",
        "  TF-Coder find solutions quickly. For instance, if you know that a reshape,\n",
        "  transpose, cast, or other similar operation should be applied to an input or\n",
        "  as the last operation to produce the output, consider applying that operation\n",
        "  manually to the input-output example, to help TF-Coder focus on the more\n",
        "  difficult parts."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n7vV5e0Iv_rz"
      },
      "source": [
        "#### Input-Output Example"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cB337V-cwDbQ"
      },
      "source": [
        "Creating a good input-output example is crucial for TF-Coder to find the\n",
        "solution you want. The example should be robust enough to rule out _false\n",
        "positive solutions_, which are TensorFlow expressions that work on the given\n",
        "example, but fail to generalize in the desired way.\n",
        "\n",
        "Here are some techniques that reduce the risk of false positives:\n",
        "\n",
        "* **Include more numbers** in the input and output tensors. TF-Coder will only\n",
        "  output a solution if it works on the provided example, so having many numbers\n",
        "  in the output tensor means it is less likely for incorrect solutions to\n",
        "  produce all of the correct numbers by chance.\n",
        "\n",
        "* **Use random-looking numbers** in the input tensors. For example,\n",
        "  `[18, 73, 34, 51]` would be a better input tensor than `[1, 2, 3, 4]`, since\n",
        "  the former is not all consecutive and not all increasing. This helps eliminate\n",
        "  patterns in the input tensors that false positive solutions can take advantage\n",
        "  of.\n",
        "\n",
        "* **Remove patterns from the output other than the intended one**. For example,\n",
        "  if the output tensor is a selection of numbers from input tensors, make sure\n",
        "  the selected numbers aren't all the maximum element along some axis, unless\n",
        "  that is the intended pattern.\n",
        "\n",
        "* **Include edge cases** where relevant. These could include negative numbers,\n",
        "  zero, or duplicate numbers, when applicable to the problem.\n",
        "\n",
        "* **Distinguish between indices and non-indices**. If you know a number should\n",
        "  not be used as an index, consider making it out of range of valid indices\n",
        "  (negative, too large, or even floating-point).\n",
        "\n",
        "* **Follow any constraints that exist in your real program**. For example, if an\n",
        "  input tensor only contains positive numbers, TF-Coder may produce a solution\n",
        "  that doesn't generalize to negative numbers. Whether this is acceptable\n",
        "  depends on whether that tensor could possibly contain negative numbers in your\n",
        "  real program. Of course, depending on the problem, a completely general\n",
        "  solution may be unnecessarily harder to find.\n",
        "\n",
        "In general, false positive solutions are more common if the output tensor\n",
        "contains a relatively low amount of information given the inputs. This may\n",
        "happen if the output is a scalar or boolean tensor, or if the output is\n",
        "constructed by selecting one or a few elements from an input. When possible, try\n",
        "to include many numbers in the output so that it contains enough information to\n",
        "unambiguously identify the intended transformation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8J9OmOq_wHFJ"
      },
      "source": [
        "#### Constants"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QfL1sYB5wJxC"
      },
      "source": [
        "* TF-Coder will print out the list of constants that it is using, including\n",
        "  constants chosen through heuristics. This list is ordered with highest-\n",
        "  priority constants at the beginning.\n",
        "* If the intended solution requires a constant that is not in TF-Coder's printed\n",
        "  list of constants, then TF-Coder will be _unable_ to find the intended\n",
        "  solution. So, it is important to provide any necessary constants.\n",
        "* If you explicitly provide constants, they will be used with the highest\n",
        "  priority. Thus, even if TF-Coder's heuristics choose your desired constant, it\n",
        "  may be better to provide the constant explicitly so that TF-Coder is more\n",
        "  confident about using your constant.\n",
        "* Providing extraneous constants will slow down the tool."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hex6HpsLwNzI"
      },
      "source": [
        "#### Description"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2GnVYA7_wL-R"
      },
      "source": [
        "* The description is optional. If provided, it is used to prioritize TensorFlow\n",
        "  operations that fit with the description.\n",
        "* If you know of a TensorFlow operation (e.g., `tf.reduce_max`) that is\n",
        "  relevant, include its name (e.g., \"tf.reduce_max\") anywhere in the\n",
        "  description. This will lead TF-Coder to prioritize that operation.\n",
        "* If possible, try to describe how the output should be computed, rather than\n",
        "  what the output conceptually represents.\n",
        "* A good description is less important than a good input-output example."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6sKleZJIwQkM"
      },
      "source": [
        "#### Other Details and Advanced Options"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WuDoERN0lx4a"
      },
      "source": [
        "* When running TF-Coder, you can set the time limit, the number of solutions to\n",
        "  find, and whether solutions are required to use inputs.\n",
        "  * Time limit: This is the maximum amount of time, in seconds, that TF-Coder\n",
        "    will spend on the problem before giving up. Note that you can stop the tool\n",
        "    at any time by pressing the cell's stop button.\n",
        "  * Number of solutions: TF-Coder can continue searching for more solutions\n",
        "    after the first solution is found. This can help you examine different ways\n",
        "    of solving the problem. However, enabling multiple solutions will cause the\n",
        "    entire search to slow down, even for the first solution.\n",
        "  * Solution requirement: By default, solutions are required to use every input\n",
        "    tensor at least once. This constraint can be relaxed to allow solutions that\n",
        "    use only one input (if there are multiple inputs), or even solutions that\n",
        "    use no inputs at all.\n",
        "\n",
        "* By default, integer tensors have a DType of `tf.int32`, and float tensors have\n",
        "  a DType of `tf.float32`. To specify a different DType, provide a `tf.Tensor`\n",
        "  object instead of a list. For example:\n",
        "  * If an input is given as `[3, 1, 7, 4]`, then it will have a DType of\n",
        "    `tf.int32`.\n",
        "  * If an input is given as `tf.constant([3, 1, 7, 4], dtype=tf.int64)`, then it\n",
        "    will have a DType of `tf.int64`.\n",
        "\n",
        "* A primitive scalar input can be specified with a Python float or int, and a\n",
        "  scalar tensor can be specified with a `tf.Tensor`:\n",
        "  * If an input is given as `[123]`, then it will be a 1-dimensional tensor with\n",
        "    shape `(1,)`, equivalent to `tf.constant([123])`.\n",
        "  * If an input is given as `123`, then it will remain a Python primitive int,\n",
        "    not a `tf.Tensor`.\n",
        "  * If an input is given as `tf.constant(123)`, then it will be a 0-dimensional\n",
        "    scalar tensor with shape `()`.\n",
        "\n",
        "* Input and output tensors can have at most 4 dimensions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pHYlFkj6l5af"
      },
      "source": [
        "## Example problems that TF-Coder can solve"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jIwb8VOombkU"
      },
      "source": [
        "Here are several examples of real-life problems that TF-Coder can solve."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QfGTQ7hF4VPV"
      },
      "source": [
        "# Real task encountered by a Googler.\n",
        "inputs = {\n",
        "    'tensor': [[0, 1, 0, 0],\n",
        "               [0, 1, 1, 0],\n",
        "               [1, 1, 1, 1]],\n",
        "}\n",
        "output = [[0.0, 1.0, 0.0, 0.0],\n",
        "          [0.0, 0.5, 0.5, 0.0],\n",
        "          [0.25, 0.25, 0.25, 0.25]]\n",
        "constants = []\n",
        "description = 'normalize the rows of a tensor'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0wzohJ299ObT"
      },
      "source": [
        "# Real task encountered by a Googler.\n",
        "inputs = {\n",
        "    'elements': [0, 0, 0, 1, 3, 3],\n",
        "}\n",
        "output = [[0, 0], [0, 1], [0, 2], [1, 0], [3, 0], [3, 1]]\n",
        "constants = []\n",
        "description = 'pair each element with a counter'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7DvZW75z4ZwN"
      },
      "source": [
        "# Real task encountered by a Googler.\n",
        "inputs = {\n",
        "    'sparse': tf.SparseTensor(\n",
        "        indices=[[0, 0, 0], [0, 1, 1], [1, 1, 1], [1, 1, 2]],\n",
        "        values=[1., 1., 1., 1.],\n",
        "        dense_shape=[2, 2, 800]),\n",
        "}\n",
        "output = tf.SparseTensor(\n",
        "    indices=[[0, 0, 0], [0, 1, 1]],\n",
        "    values=[1., 1.],\n",
        "    dense_shape=[1, 2, 800])\n",
        "constants = []\n",
        "description = 'slice index 0 of the first dimension of a SparseTensor'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rv8jFiQvQYEC"
      },
      "source": [
        "# Real task encountered by a Googler.\n",
        "inputs = {\n",
        "    'lengths': [3, 4, 2, 1],\n",
        "}\n",
        "output = [[1, 1, 1, 0, 0],\n",
        "          [1, 1, 1, 1, 0],\n",
        "          [1, 1, 0, 0, 0],\n",
        "          [1, 0, 0, 0, 0]]\n",
        "constants = [5]\n",
        "description = 'create a mask for sequences of the given lengths'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pL50uFHLyu8l"
      },
      "source": [
        "# Real task encountered by a Googler.\n",
        "inputs = {\n",
        "    'segments': [ 1,  1,  1,  0,  0,  2],\n",
        "    'data':     [10, 20, 30, 14, 15, 26],\n",
        "}\n",
        "output = [14, 15, 10, 20, 30, 26]\n",
        "constants = []\n",
        "description = 'sort the segments'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DCvkqaXs4fyU"
      },
      "source": [
        "# Adapted from https://stackoverflow.com/questions/53054668\n",
        "inputs = {\n",
        "    'values': [37, 42, 42, 37, 28, 15, 42, 15],\n",
        "}\n",
        "output = [0, 1, 1, 0, 2, 3, 1, 3]\n",
        "constants = []\n",
        "description = 'group items by value and get the group indices'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A4hI5vMLYQNj"
      },
      "source": [
        "# Adapted from https://stackoverflow.com/questions/47816231\n",
        "inputs = {\n",
        "    'vector': [3, 5, 0, 2, 3, 3, 0],\n",
        "}\n",
        "output = [[1., 0., 0., 0., 1., 1., 0.],\n",
        "          [0., 1., 0., 0., 0., 0., 0.],\n",
        "          [0., 0., 1., 0., 0., 0., 1.],\n",
        "          [0., 0., 0., 1., 0., 0., 0.],\n",
        "          [1., 0., 0., 0., 1., 1., 0.],\n",
        "          [1., 0., 0., 0., 1., 1., 0.],\n",
        "          [0., 0., 1., 0., 0., 0., 1.]]\n",
        "constants = []\n",
        "description = 'binary tensor from vector indicating if elements are equal'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TbbocD0o0PYC"
      },
      "source": [
        "# Adapted from https://stackoverflow.com/questions/44834739\n",
        "inputs = {\n",
        "    'scores': [[0.7, 0.2, 0.1],\n",
        "               [0.4, 0.5, 0.1],\n",
        "               [0.4, 0.4, 0.2],\n",
        "               [0.3, 0.4, 0.3],\n",
        "               [0.0, 0.0, 1.0]],\n",
        "}\n",
        "output = [[1, 0, 0],\n",
        "          [0, 1, 0],\n",
        "          [1, 0, 0],\n",
        "          [0, 1, 0],\n",
        "          [0, 0, 1]]\n",
        "constants = []\n",
        "description = 'compute argmax in each tensor and set it to 1'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FUUUBi7HI3Hr"
      },
      "source": [
        "# Adapted from https://stackoverflow.com/questions/33769041\n",
        "inputs = {\n",
        "    'first': [-1, 0, -3, 2, 1, 3, 5, -1, -9, 2, 10],\n",
        "    'second': [12, 3, 45, 6, 7, 8, 9, 87, 65, 4, 32],\n",
        "}\n",
        "output = [6, 8, 9, 4, 32]\n",
        "constants = [1]\n",
        "description = 'select the values in the second tensor where the first tensor is greater than 1'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q6uRr4x9WHRC"
      },
      "source": [
        "## Supported Operations"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XvHDbukUWNKn",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "80d91e78-31a0-45d2-d4e4-692a4064e3e6"
      },
      "source": [
        "# Run this cell to print all supported operations.\n",
        "colab_interface.print_supported_operations()"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "TensorFlow functions:\n",
            "---------------------\n",
            "tf.abs(x)\n",
            "tf.add(x, y)\n",
            "tf.add_n(inputs)\n",
            "tf.argmax(input, axis)\n",
            "tf.argmin(input, axis)\n",
            "tf.argsort(values, axis, stable=True)\n",
            "tf.argsort(values, axis, direction='DESCENDING', stable=True)\n",
            "tf.boolean_mask(tensor, mask)\n",
            "tf.broadcast_to(input, shape)\n",
            "tf.cast(x, dtype)\n",
            "tf.clip_by_value(t, clip_value_min, clip_value_max)\n",
            "tf.concat(values, axis)\n",
            "tf.constant(value)\n",
            "tf.constant(value, dtype)\n",
            "tf.divide(x, y)\n",
            "tf.equal(x, y)\n",
            "tf.exp(x)\n",
            "tf.expand_dims(input, axis)\n",
            "tf.eye(num_rows)\n",
            "tf.eye(num_rows, num_columns)\n",
            "tf.eye(num_rows, dtype)\n",
            "tf.fill(dims, value)\n",
            "tf.gather(params, indices)\n",
            "tf.gather(params, indices, axis, batch_dims)\n",
            "tf.gather_nd(params, indices)\n",
            "tf.gather_nd(params, indices, batch_dims)\n",
            "tf.greater(x, y)\n",
            "tf.greater_equal(x, y)\n",
            "tf.math.bincount(arr)\n",
            "tf.math.ceil(x)\n",
            "tf.math.count_nonzero(input)\n",
            "tf.math.count_nonzero(input, axis)\n",
            "tf.math.cumsum(x, axis)\n",
            "tf.math.cumsum(x, axis, exclusive=True)\n",
            "tf.math.divide_no_nan(x, y)\n",
            "tf.math.floor(x)\n",
            "tf.math.log(x)\n",
            "tf.math.negative(x)\n",
            "tf.math.reciprocal(x)\n",
            "tf.math.reciprocal_no_nan(x)\n",
            "tf.math.segment_max(data, segment_ids)\n",
            "tf.math.segment_mean(data, segment_ids)\n",
            "tf.math.segment_min(data, segment_ids)\n",
            "tf.math.segment_prod(data, segment_ids)\n",
            "tf.math.segment_sum(data, segment_ids)\n",
            "tf.math.squared_difference(x, y)\n",
            "tf.math.top_k(input, k)\n",
            "tf.math.unsorted_segment_max(data, segment_ids, num_segments)\n",
            "tf.math.unsorted_segment_mean(data, segment_ids, num_segments)\n",
            "tf.math.unsorted_segment_min(data, segment_ids, num_segments)\n",
            "tf.math.unsorted_segment_prod(data, segment_ids, num_segments)\n",
            "tf.math.unsorted_segment_sum(data, segment_ids, num_segments)\n",
            "tf.matmul(a, b)\n",
            "tf.maximum(x, y)\n",
            "tf.minimum(x, y)\n",
            "tf.multiply(x, y)\n",
            "tf.not_equal(x, y)\n",
            "tf.one_hot(indices, depth)\n",
            "tf.ones(shape)\n",
            "tf.ones_like(input)\n",
            "tf.pad(tensor, paddings, mode='CONSTANT')\n",
            "tf.pad(tensor, paddings, mode='CONSTANT', constant_values)\n",
            "tf.pad(tensor, paddings, mode='REFLECT')\n",
            "tf.pad(tensor, paddings, mode='SYMMETRIC')\n",
            "tf.range(start)\n",
            "tf.range(start, limit, delta)\n",
            "tf.reduce_any(input_tensor, axis)\n",
            "tf.reduce_max(input_tensor)\n",
            "tf.reduce_max(input_tensor, axis)\n",
            "tf.reduce_mean(input_tensor)\n",
            "tf.reduce_mean(input_tensor, axis)\n",
            "tf.reduce_min(input_tensor)\n",
            "tf.reduce_min(input_tensor, axis)\n",
            "tf.reduce_prod(input_tensor, axis)\n",
            "tf.reduce_sum(input_tensor)\n",
            "tf.reduce_sum(input_tensor, axis)\n",
            "tf.reshape(tensor, shape)\n",
            "tf.reverse(tensor, axis)\n",
            "tf.roll(input, shift, axis)\n",
            "tf.round(x)\n",
            "tf.scatter_nd(indices, updates, shape)\n",
            "tf.searchsorted(sorted_sequence, values, side='left')\n",
            "tf.searchsorted(sorted_sequence, values, side='right')\n",
            "tf.sequence_mask(lengths)\n",
            "tf.sequence_mask(lengths, maxlen)\n",
            "tf.shape(input)\n",
            "tf.sign(x)\n",
            "tf.sort(values, axis)\n",
            "tf.sort(values, axis, direction='DESCENDING')\n",
            "tf.sqrt(x)\n",
            "tf.square(x)\n",
            "tf.squeeze(input)\n",
            "tf.squeeze(input, axis)\n",
            "tf.stack(values, axis)\n",
            "tf.subtract(x, y)\n",
            "tf.tensor_scatter_nd_update(tensor, indices, updates)\n",
            "tf.tensordot(a, b, axes)\n",
            "tf.tile(input, multiples)\n",
            "tf.transpose(a)\n",
            "tf.transpose(a, perm)\n",
            "tf.unique_with_counts(x)\n",
            "tf.unstack(value, axis)\n",
            "tf.where(condition)\n",
            "tf.where(condition, x, y)\n",
            "tf.zeros(shape)\n",
            "tf.zeros_like(input)\n",
            "\n",
            "SparseTensor functions:\n",
            "-----------------------\n",
            "tf.SparseTensor(indices, values, dense_shape)\n",
            "tf.sparse.add(a, b)\n",
            "tf.sparse.concat(axis, sp_inputs)\n",
            "tf.sparse.expand_dims(sp_input, axis)\n",
            "tf.sparse.from_dense(tensor)\n",
            "tf.sparse.maximum(sp_a, sp_b)\n",
            "tf.sparse.minimum(sp_a, sp_b)\n",
            "tf.sparse.reduce_max(sp_input, axis, output_is_sparse)\n",
            "tf.sparse.reduce_sum(sp_input, axis, output_is_sparse)\n",
            "tf.sparse.reset_shape(sp_input)\n",
            "tf.sparse.reshape(sp_input, shape)\n",
            "tf.sparse.retain(sp_input, to_retain)\n",
            "tf.sparse.slice(sp_input, start, size)\n",
            "tf.sparse.split(sp_input, num_split, axis)\n",
            "tf.sparse.to_dense(sp_input)\n",
            "tf.sparse.to_dense(sp_input, default_value)\n",
            "tf.sparse.to_indicator(sp_input, vocab_size)\n",
            "tf.sparse.transpose(sp_input)\n",
            "tf.sparse.transpose(sp_input, perm)\n",
            "\n",
            "Python-syntax operations:\n",
            "-------------------------\n",
            "IndexingAxis1Operation:             arg1[:, arg2]\n",
            "IndexingOperation:                  arg1[arg2]\n",
            "PairCreationOperation:              (arg1, arg2)\n",
            "SingletonTupleCreationOperation:    (arg1,)\n",
            "SlicingAxis0BothOperation:          arg1[arg2:arg3]\n",
            "SlicingAxis0LeftOperation:          arg1[arg2:]\n",
            "SlicingAxis0RightOperation:         arg1[:arg2]\n",
            "SlicingAxis1BothOperation:          arg1[:, arg2:arg3]\n",
            "SlicingAxis1LeftOperation:          arg1[:, arg2:]\n",
            "SlicingAxis1RightOperation:         arg1[:, :arg2]\n",
            "TripleCreationOperation:            (arg1, arg2, arg3)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NYWv_mSbmfOO"
      },
      "source": [
        "## Feedback? Questions?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NCdv8OzQmn0j"
      },
      "source": [
        "More information and resources about TF-Coder can be found at our [GitHub repo](https://github.com/google-research/tensorflow-coder).\n",
        "\n",
        "To report a bug or make a feature request, please raise a\n",
        "[GitHub issue](https://github.com/google-research/tensorflow-coder/issues).\n",
        "\n",
        "If you have accidentally run the TF-Coder tool on *personal or confidential information*, and you have agreed to the *public release* of your TF-Coder usage data, you may reach out to tf-coder-support@google.com to request removal of your data from the public release. Such requests will be handled on a best-effort basis with no guarantee of success. Again, please do not run the TF-Coder tool on personal or confidential information.\n",
        "\n",
        "This is a research project, not an official Google product."
      ]
    }
  ]
}