Skip to content

Instantly share code, notes, and snippets.

@gpantalos
Last active November 7, 2021 00:15
Show Gist options
  • Select an option

  • Save gpantalos/cbec955549f22d8dbbcc41fa4dfb61ad to your computer and use it in GitHub Desktop.

Select an option

Save gpantalos/cbec955549f22d8dbbcc41fa4dfb61ad to your computer and use it in GitHub Desktop.

Revisions

  1. Georges Pantalos revised this gist Nov 7, 2021. No changes.
  2. Georges Pantalos revised this gist Feb 18, 2021. 1 changed file with 31 additions and 20 deletions.
    51 changes: 31 additions & 20 deletions stein_identity.ipynb
    Original file line number Diff line number Diff line change
    @@ -6,7 +6,7 @@
    "name": "stein_identity.ipynb",
    "provenance": [],
    "collapsed_sections": [],
    "authorship_tag": "ABX9TyNB8oQc1T3c7sxMsXPW2CLH",
    "authorship_tag": "ABX9TyM+COvZLzmhu/hW65ZSfSP4",
    "include_colab_link": true
    },
    "kernelspec": {
    @@ -50,7 +50,7 @@
    "tfd = tfp.distributions\n",
    "tf.random.set_seed(0)"
    ],
    "execution_count": null,
    "execution_count": 2,
    "outputs": []
    },
    {
    @@ -69,13 +69,13 @@
    "base_uri": "https://localhost:8080/"
    },
    "id": "1JM92iB7VDKq",
    "outputId": "11a8c920-3f40-43b9-9ad2-78e589afc953"
    "outputId": "65abf331-b2db-4730-8b32-280305a98781"
    },
    "source": [
    "p = tfd.Normal(0,1)\n",
    "p"
    ],
    "execution_count": null,
    "execution_count": 3,
    "outputs": [
    {
    "output_type": "execute_result",
    @@ -87,7 +87,7 @@
    "metadata": {
    "tags": []
    },
    "execution_count": 61
    "execution_count": 3
    }
    ]
    },
    @@ -107,14 +107,14 @@
    "base_uri": "https://localhost:8080/"
    },
    "id": "pO5icgaEVG3p",
    "outputId": "7b36f2f9-1a9a-4411-a20f-c1740679dea5"
    "outputId": "2351337b-2c1a-4e7f-c545-f3748ded9481"
    },
    "source": [
    "x = p.sample(1e5)\n",
    "x = tf.Variable(x)\n",
    "x.shape"
    ],
    "execution_count": null,
    "execution_count": 4,
    "outputs": [
    {
    "output_type": "execute_result",
    @@ -126,7 +126,7 @@
    "metadata": {
    "tags": []
    },
    "execution_count": 62
    "execution_count": 4
    }
    ]
    },
    @@ -146,15 +146,15 @@
    "base_uri": "https://localhost:8080/"
    },
    "id": "uzkWCSh0Wr59",
    "outputId": "380ea18d-e537-4880-ba10-777e084389fc"
    "outputId": "c1a470f8-d761-4286-87f5-1006ffb24183"
    },
    "source": [
    "with tf.GradientTape() as tape:\n",
    " logp = p.log_prob(x)\n",
    "score = tape.gradient(logp, x)\n",
    "score.shape"
    ],
    "execution_count": null,
    "execution_count": 5,
    "outputs": [
    {
    "output_type": "execute_result",
    @@ -166,7 +166,7 @@
    "metadata": {
    "tags": []
    },
    "execution_count": 63
    "execution_count": 5
    }
    ]
    },
    @@ -194,7 +194,7 @@
    " # add smooth functions...\n",
    "]"
    ],
    "execution_count": null,
    "execution_count": 6,
    "outputs": []
    },
    {
    @@ -213,7 +213,7 @@
    "base_uri": "https://localhost:8080/"
    },
    "id": "nH6RfK_bWWYP",
    "outputId": "abf5272d-7424-4919-c4c3-6f19ee87b1e0"
    "outputId": "7c1301bb-357a-4de3-827a-6ed74b4bec13"
    },
    "source": [
    "grad_phi_x = []\n",
    @@ -230,7 +230,7 @@
    "grad_phi_x = tf.stack(grad_phi_x)\n",
    "grad_phi_x.shape"
    ],
    "execution_count": null,
    "execution_count": 7,
    "outputs": [
    {
    "output_type": "execute_result",
    @@ -242,7 +242,7 @@
    "metadata": {
    "tags": []
    },
    "execution_count": 65
    "execution_count": 7
    }
    ]
    },
    @@ -262,24 +262,24 @@
    "base_uri": "https://localhost:8080/"
    },
    "id": "T8flr6qdXxFn",
    "outputId": "84575c96-0ec6-45d7-dcc6-963d545746b9"
    "outputId": "261b0cc5-f333-4656-8ba2-5fa12aaa7004"
    },
    "source": [
    "tf.reduce_mean(phi_xs * score + grad_phi_x).numpy()"
    "tf.reduce_mean(phi_xs * score + grad_phi_x).numpy() ** 2"
    ],
    "execution_count": null,
    "execution_count": 8,
    "outputs": [
    {
    "output_type": "execute_result",
    "data": {
    "text/plain": [
    "-0.004825321"
    "5.552772948637223e-05"
    ]
    },
    "metadata": {
    "tags": []
    },
    "execution_count": 70
    "execution_count": 8
    }
    ]
    },
    @@ -291,6 +291,17 @@
    "source": [
    "Indeed approaches 0!"
    ]
    },
    {
    "cell_type": "code",
    "metadata": {
    "id": "c7X61SRtvrCQ"
    },
    "source": [
    ""
    ],
    "execution_count": 8,
    "outputs": []
    }
    ]
    }
  3. Georges Pantalos revised this gist Feb 18, 2021. 1 changed file with 9 additions and 9 deletions.
    18 changes: 9 additions & 9 deletions stein_identity.ipynb
    Original file line number Diff line number Diff line change
    @@ -6,7 +6,7 @@
    "name": "stein_identity.ipynb",
    "provenance": [],
    "collapsed_sections": [],
    "authorship_tag": "ABX9TyN75IHlGoQnMK/MdW/SPfkc",
    "authorship_tag": "ABX9TyNB8oQc1T3c7sxMsXPW2CLH",
    "include_colab_link": true
    },
    "kernelspec": {
    @@ -50,7 +50,7 @@
    "tfd = tfp.distributions\n",
    "tf.random.set_seed(0)"
    ],
    "execution_count": 1,
    "execution_count": null,
    "outputs": []
    },
    {
    @@ -59,7 +59,7 @@
    "id": "dny1VJB9V3nL"
    },
    "source": [
    "Define $p = \\mathcal N(0,1)$"
    "Define a distribution $p = \\mathcal N(0,1)$"
    ]
    },
    {
    @@ -75,7 +75,7 @@
    "p = tfd.Normal(0,1)\n",
    "p"
    ],
    "execution_count": 61,
    "execution_count": null,
    "outputs": [
    {
    "output_type": "execute_result",
    @@ -114,7 +114,7 @@
    "x = tf.Variable(x)\n",
    "x.shape"
    ],
    "execution_count": 62,
    "execution_count": null,
    "outputs": [
    {
    "output_type": "execute_result",
    @@ -154,7 +154,7 @@
    "score = tape.gradient(logp, x)\n",
    "score.shape"
    ],
    "execution_count": 63,
    "execution_count": null,
    "outputs": [
    {
    "output_type": "execute_result",
    @@ -194,7 +194,7 @@
    " # add smooth functions...\n",
    "]"
    ],
    "execution_count": 64,
    "execution_count": null,
    "outputs": []
    },
    {
    @@ -230,7 +230,7 @@
    "grad_phi_x = tf.stack(grad_phi_x)\n",
    "grad_phi_x.shape"
    ],
    "execution_count": 65,
    "execution_count": null,
    "outputs": [
    {
    "output_type": "execute_result",
    @@ -267,7 +267,7 @@
    "source": [
    "tf.reduce_mean(phi_xs * score + grad_phi_x).numpy()"
    ],
    "execution_count": 70,
    "execution_count": null,
    "outputs": [
    {
    "output_type": "execute_result",
  4. Georges Pantalos revised this gist Feb 18, 2021. No changes.
  5. Georges Pantalos revised this gist Feb 18, 2021. 1 changed file with 12 additions and 1 deletion.
    13 changes: 12 additions & 1 deletion stein_identity.ipynb
    Original file line number Diff line number Diff line change
    @@ -6,14 +6,25 @@
    "name": "stein_identity.ipynb",
    "provenance": [],
    "collapsed_sections": [],
    "authorship_tag": "ABX9TyN75IHlGoQnMK/MdW/SPfkc"
    "authorship_tag": "ABX9TyN75IHlGoQnMK/MdW/SPfkc",
    "include_colab_link": true
    },
    "kernelspec": {
    "name": "python3",
    "display_name": "Python 3"
    }
    },
    "cells": [
    {
    "cell_type": "markdown",
    "metadata": {
    "id": "view-in-github",
    "colab_type": "text"
    },
    "source": [
    "<a href=\"https://colab.research.google.com/gist/gpantalos/cbec955549f22d8dbbcc41fa4dfb61ad/stein_identity.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
    ]
    },
    {
    "cell_type": "markdown",
    "metadata": {
  6. Georges Pantalos created this gist Feb 18, 2021.
    285 changes: 285 additions & 0 deletions stein_identity.ipynb
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,285 @@
    {
    "nbformat": 4,
    "nbformat_minor": 0,
    "metadata": {
    "colab": {
    "name": "stein_identity.ipynb",
    "provenance": [],
    "collapsed_sections": [],
    "authorship_tag": "ABX9TyN75IHlGoQnMK/MdW/SPfkc"
    },
    "kernelspec": {
    "name": "python3",
    "display_name": "Python 3"
    }
    },
    "cells": [
    {
    "cell_type": "markdown",
    "metadata": {
    "id": "HjL76oTdVP_M"
    },
    "source": [
    "# Stein's Identity in TensorFlow\n",
    "This is Stein's identity\n",
    "$$\n",
    "\\mathbb{E}_{x \\sim p}\\left[\\phi(x) \\nabla_{x} \\log p(x)^{\\top}+\\nabla_{x} \\phi(x)\\right]=0\n",
    "$$\n",
    "Let's test it in TF."
    ]
    },
    {
    "cell_type": "code",
    "metadata": {
    "id": "3BTuVKAxU62o"
    },
    "source": [
    "import tensorflow as tf\n",
    "import tensorflow_probability as tfp\n",
    "tfd = tfp.distributions\n",
    "tf.random.set_seed(0)"
    ],
    "execution_count": 1,
    "outputs": []
    },
    {
    "cell_type": "markdown",
    "metadata": {
    "id": "dny1VJB9V3nL"
    },
    "source": [
    "Define $p = \\mathcal N(0,1)$"
    ]
    },
    {
    "cell_type": "code",
    "metadata": {
    "colab": {
    "base_uri": "https://localhost:8080/"
    },
    "id": "1JM92iB7VDKq",
    "outputId": "11a8c920-3f40-43b9-9ad2-78e589afc953"
    },
    "source": [
    "p = tfd.Normal(0,1)\n",
    "p"
    ],
    "execution_count": 61,
    "outputs": [
    {
    "output_type": "execute_result",
    "data": {
    "text/plain": [
    "<tfp.distributions.Normal 'Normal' batch_shape=[] event_shape=[] dtype=float32>"
    ]
    },
    "metadata": {
    "tags": []
    },
    "execution_count": 61
    }
    ]
    },
    {
    "cell_type": "markdown",
    "metadata": {
    "id": "gKNYVgbxV-iI"
    },
    "source": [
    "Sample $x \\sim p$"
    ]
    },
    {
    "cell_type": "code",
    "metadata": {
    "colab": {
    "base_uri": "https://localhost:8080/"
    },
    "id": "pO5icgaEVG3p",
    "outputId": "7b36f2f9-1a9a-4411-a20f-c1740679dea5"
    },
    "source": [
    "x = p.sample(1e5)\n",
    "x = tf.Variable(x)\n",
    "x.shape"
    ],
    "execution_count": 62,
    "outputs": [
    {
    "output_type": "execute_result",
    "data": {
    "text/plain": [
    "TensorShape([100000])"
    ]
    },
    "metadata": {
    "tags": []
    },
    "execution_count": 62
    }
    ]
    },
    {
    "cell_type": "markdown",
    "metadata": {
    "id": "5qAhWPeiWpUZ"
    },
    "source": [
    "Compute and evaluate score function $ \\nabla_{x} \\log p(x)$"
    ]
    },
    {
    "cell_type": "code",
    "metadata": {
    "colab": {
    "base_uri": "https://localhost:8080/"
    },
    "id": "uzkWCSh0Wr59",
    "outputId": "380ea18d-e537-4880-ba10-777e084389fc"
    },
    "source": [
    "with tf.GradientTape() as tape:\n",
    " logp = p.log_prob(x)\n",
    "score = tape.gradient(logp, x)\n",
    "score.shape"
    ],
    "execution_count": 63,
    "outputs": [
    {
    "output_type": "execute_result",
    "data": {
    "text/plain": [
    "TensorShape([100000])"
    ]
    },
    "metadata": {
    "tags": []
    },
    "execution_count": 63
    }
    ]
    },
    {
    "cell_type": "markdown",
    "metadata": {
    "id": "ZMP3J2HTVuus"
    },
    "source": [
    "Define a mapping $\\phi(x)$ of any number of smooth functions"
    ]
    },
    {
    "cell_type": "code",
    "metadata": {
    "id": "UsvaLrEUVKLa"
    },
    "source": [
    "phi = [\n",
    " lambda i: i, \n",
    " lambda i: i ** 2,\n",
    " lambda i: tf.exp(i), \n",
    " lambda i: tf.sin(i), \n",
    " \n",
    " # add smooth functions...\n",
    "]"
    ],
    "execution_count": 64,
    "outputs": []
    },
    {
    "cell_type": "markdown",
    "metadata": {
    "id": "I-9hirvqXXsL"
    },
    "source": [
    "Compute gradient wrt inputs $\\nabla_{x} \\phi(x)$"
    ]
    },
    {
    "cell_type": "code",
    "metadata": {
    "colab": {
    "base_uri": "https://localhost:8080/"
    },
    "id": "nH6RfK_bWWYP",
    "outputId": "abf5272d-7424-4919-c4c3-6f19ee87b1e0"
    },
    "source": [
    "grad_phi_x = []\n",
    "phi_xs = []\n",
    "\n",
    "for mapping in phi:\n",
    " with tf.GradientTape() as tape:\n",
    " phi_x = mapping(x) \n",
    " phi_xs.append(phi_x)\n",
    " grad_phi_x.append(tape.gradient(phi_x, x))\n",
    "\n",
    "# stack gradients\n",
    "phi_xs = tf.stack(phi_xs)\n",
    "grad_phi_x = tf.stack(grad_phi_x)\n",
    "grad_phi_x.shape"
    ],
    "execution_count": 65,
    "outputs": [
    {
    "output_type": "execute_result",
    "data": {
    "text/plain": [
    "TensorShape([4, 100000])"
    ]
    },
    "metadata": {
    "tags": []
    },
    "execution_count": 65
    }
    ]
    },
    {
    "cell_type": "markdown",
    "metadata": {
    "id": "ONnBztPCX1uc"
    },
    "source": [
    "Combine previous results to verify Stein's identity"
    ]
    },
    {
    "cell_type": "code",
    "metadata": {
    "colab": {
    "base_uri": "https://localhost:8080/"
    },
    "id": "T8flr6qdXxFn",
    "outputId": "84575c96-0ec6-45d7-dcc6-963d545746b9"
    },
    "source": [
    "tf.reduce_mean(phi_xs * score + grad_phi_x).numpy()"
    ],
    "execution_count": 70,
    "outputs": [
    {
    "output_type": "execute_result",
    "data": {
    "text/plain": [
    "-0.004825321"
    ]
    },
    "metadata": {
    "tags": []
    },
    "execution_count": 70
    }
    ]
    },
    {
    "cell_type": "markdown",
    "metadata": {
    "id": "ijNbFJlmYukL"
    },
    "source": [
    "Indeed approaches 0!"
    ]
    }
    ]
    }