{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "JAX Colab GPU Test",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/google/jax/blob/master/tests/notebooks/colab_gpu.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WkadOyTDCAWD",
        "colab_type": "text"
      },
      "source": [
        "# JAX Colab GPU Test\n",
        "\n",
        "This notebook is meant to be run in a [Colab](http://colab.research.google.com) GPU runtime as a basic check for JAX updates."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_tKNrbqqBHwu",
        "colab_type": "code",
        "outputId": "ae4a051a-91ed-4742-c8e1-31de8304ef33",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "import jax\n",
        "import jaxlib\n",
        "\n",
        "!cat /var/colab/hostname\n",
        "print(jax.__version__)\n",
        "print(jaxlib.__version__)"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "gpu-t4-s-kbefivsjoreh\n",
            "0.1.64\n",
            "0.1.45\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oqEG21rADO1F",
        "colab_type": "text"
      },
      "source": [
        "## Confirm Device"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "8BwzMYhKGQj6",
        "outputId": "ff4f52b3-f7bb-468a-c1ad-debe65841f3f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "from jaxlib import xla_extension\n",
        "import jax\n",
        "key = jax.random.PRNGKey(1701)\n",
        "arr = jax.random.normal(key, (1000,))\n",
        "device = arr.device_buffer.device()\n",
        "print(f\"JAX device type: {device}\")\n",
        "assert isinstance(device, xla_extension.GpuDevice), \"unexpected JAX device type\""
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "JAX device type: gpu:0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z0FUY9yUC4k1",
        "colab_type": "text"
      },
      "source": [
        "## Matrix Multiplication"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "eXn8GUl6CG5N",
        "outputId": "688c37f3-e830-4ba8-b1e6-b4e014cb11a9",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "import jax\n",
        "import numpy as np\n",
        "\n",
        "# matrix multiplication on GPU\n",
        "key = jax.random.PRNGKey(0)\n",
        "x = jax.random.normal(key, (3000, 3000))\n",
        "result = jax.numpy.dot(x, x.T).mean()\n",
        "print(result)"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "1.0216676\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0zTA2Q19DW4G",
        "colab_type": "text"
      },
      "source": [
        "## Linear Algebra"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uW9j84_UDYof",
        "colab_type": "code",
        "outputId": "80069760-12ab-4df2-9f5c-be2536de59b7",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "import jax.numpy as jnp\n",
        "import jax.random as rand\n",
        "\n",
        "N = 10\n",
        "M = 20\n",
        "key = rand.PRNGKey(1701)\n",
        "\n",
        "X = rand.normal(key, (N, M))\n",
        "u, s, vt = jnp.linalg.svd(X)\n",
        "assert u.shape == (N, N)\n",
        "assert vt.shape == (M, M)\n",
        "print(s)"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[6.9178247 5.9580336 5.5811076 4.5069666 4.1115823 3.9735446 3.3307252\n",
            " 2.866489  1.8229384 1.5478926]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jCyKUn4-DCXn",
        "colab_type": "text"
      },
      "source": [
        "## XLA Compilation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "2GOn_HhDPuEn",
        "outputId": "a51d7d07-8513-4503-bceb-d5b0e2b4e4a8",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "@jax.jit\n",
        "def selu(x, alpha=1.67, lmbda=1.05):\n",
        "  return lmbda * jax.numpy.where(x > 0, x, alpha * jax.numpy.exp(x) - alpha)\n",
        "x = jax.random.normal(key, (5000,))\n",
        "result = selu(x).block_until_ready()\n",
        "print(result)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[ 0.34676838 -0.7532232   1.7060698  ...  2.1208055  -0.42621925\n",
            "  0.13093245]\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}