Starting a colab to demonstrate Trax layers.

PiperOrigin-RevId: 259802336

Author: Lukasz Kaiser

tensor2tensor/trax/layers/demo.ipynb

@@ -0,0 +1,189 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "7yuytuIllsv1"
+      },
+      "source": [
+        "# Trax Layers\n",
+        "\n",
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "colab": {
+          "height": 119
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 88481,
+          "status": "ok",
+          "timestamp": 1563927238895,
+          "user": {
+            "displayName": "Lukasz Kaiser",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mC8pChl87HbK_eOtVhtNPwUVx8btvfyYzH9UHn3=s64",
+            "userId": "13267693649565518272"
+          },
+          "user_tz": 420
+        },
+        "id": "oILRLCWN_16u",
+        "outputId": "3f750014-c633-4162-ad07-f3c56c273304"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "T2T: skipped importing 1 data_generators modules. OK if no other errors. Depend on _heavy or problem-specific py_binary targets if trying to use a module that was skipped.\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "matplotlib.use() must be called *before* pylab, matplotlib.pyplot,\n",
+            "or matplotlib.backends is imported for the first time.\n",
+            "\n"
+          ]
+        }
+      ],
+      "source": [
+        "#@title Licence and python imports.\n",
+        "# Copyright 2018 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.\n",
+        "\n",
+        "import datetime\n",
+        "import numpy as onp\n",
+        "\n",
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 5,
+      "metadata": {
+        "cellView": "both",
+        "colab": {
+          "height": 51
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 2992,
+          "status": "ok",
+          "timestamp": 1563927313403,
+          "user": {
+            "displayName": "Lukasz Kaiser",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mC8pChl87HbK_eOtVhtNPwUVx8btvfyYzH9UHn3=s64",
+            "userId": "13267693649565518272"
+          },
+          "user_tz": 420
+        },
+        "id": "vlGjGoGMTt-D",
+        "outputId": "7a7b5a1e-c01e-4a5e-eeb1-a88d9500aad4"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "/bin/sh: pip: command not found\n",
+            "/bin/sh: pip: command not found\n"
+          ]
+        }
+      ],
+      "source": [
+        "#@title Install TensorFlow, Tensor2Tensor and Trax.\n",
+        "\n",
+        "! pip install -q tensorflow\n",
+        "! pip install -q -U tensor2tensor\n",
+        "\n",
+        "from tensor2tensor.trax import trax\n",
+        "from tensor2tensor.trax import backend\n",
+        "from tensor2tensor.trax import layers as tl"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 6,
+      "metadata": {
+        "colab": {
+          "height": 68
+        },
+        "colab_type": "code",
+        "executionInfo": {
+          "elapsed": 434,
+          "status": "ok",
+          "timestamp": 1563927323194,
+          "user": {
+            "displayName": "Lukasz Kaiser",
+            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mC8pChl87HbK_eOtVhtNPwUVx8btvfyYzH9UHn3=s64",
+            "userId": "13267693649565518272"
+          },
+          "user_tz": 420
+        },
+        "id": "V09viOSEQvQe",
+        "outputId": "822ec44b-1e2e-4b6e-9fc7-9fc29f5d3783"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "[[-7 -6 -5 -4 -3]\n",
+            " [-2 -1  0  1  2]\n",
+            " [ 3  4  5  6  7]]\n",
+            "[[0 0 0 0 0]\n",
+            " [0 0 0 1 2]\n",
+            " [3 4 5 6 7]]\n"
+          ]
+        }
+      ],
+      "source": [
+        "x  = onp.arange(-7, 8).reshape(3, -1)\n",
+        "rng = backend.random.get_prng(0)\n",
+        "layer = tl.Relu()\n",
+        "params = layer.initialize(x.shape, x.dtype, rng)\n",
+        "output = layer(x, params, rng=rng)\n",
+        "print(x)\n",
+        "print(output)"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "last_runtime": {
+        "build_target": "//learning/deepmind/dm_python:dm_notebook",
+        "kind": "private"
+      },
+      "name": "Trax Layers",
+      "provenance": [
+        {
+          "file_id": "1EH76AWQ_pvT4i8ZXfkv-SCV4MrmllEl5",
+          "timestamp": 1563927451951
+        }
+      ],
+      "version": "0.3.2"
+    },
+    "kernelspec": {
+      "display_name": "Python 2",
+      "name": "python2"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}