Last active
November 7, 2021 00:15
-
-
Save gpantalos/cbec955549f22d8dbbcc41fa4dfb61ad to your computer and use it in GitHub Desktop.
Revisions
-
Georges Pantalos revised this gist
Nov 7, 2021 . No changes.There are no files selected for viewing
-
Georges Pantalos revised this gist
Feb 18, 2021 . 1 changed file with 31 additions and 20 deletions.There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode charactersOriginal file line number Diff line number Diff line change @@ -6,7 +6,7 @@ "name": "stein_identity.ipynb", "provenance": [], "collapsed_sections": [], "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": 2, "outputs": [] }, { @@ -69,13 +69,13 @@ "base_uri": "https://localhost:8080/" }, "id": "1JM92iB7VDKq", "outputId": "65abf331-b2db-4730-8b32-280305a98781" }, "source": [ "p = tfd.Normal(0,1)\n", "p" ], "execution_count": 3, "outputs": [ { "output_type": "execute_result", @@ -87,7 +87,7 @@ "metadata": { "tags": [] }, "execution_count": 3 } ] }, @@ -107,14 +107,14 @@ "base_uri": "https://localhost:8080/" }, "id": "pO5icgaEVG3p", "outputId": "2351337b-2c1a-4e7f-c545-f3748ded9481" }, "source": [ "x = p.sample(1e5)\n", "x = tf.Variable(x)\n", "x.shape" ], "execution_count": 4, "outputs": [ { "output_type": "execute_result", @@ -126,7 +126,7 @@ "metadata": { "tags": [] }, "execution_count": 4 } ] }, @@ -146,15 +146,15 @@ "base_uri": "https://localhost:8080/" }, "id": "uzkWCSh0Wr59", "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": 5, "outputs": [ { "output_type": "execute_result", @@ -166,7 +166,7 @@ "metadata": { "tags": [] }, "execution_count": 5 } ] }, @@ -194,7 +194,7 @@ " # add smooth functions...\n", "]" ], "execution_count": 6, "outputs": [] }, { @@ -213,7 +213,7 @@ "base_uri": "https://localhost:8080/" }, "id": "nH6RfK_bWWYP", "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": 7, "outputs": [ { "output_type": "execute_result", @@ -242,7 +242,7 @@ "metadata": { "tags": [] }, "execution_count": 7 } ] }, @@ -262,24 +262,24 @@ "base_uri": "https://localhost:8080/" }, "id": "T8flr6qdXxFn", "outputId": "261b0cc5-f333-4656-8ba2-5fa12aaa7004" }, "source": [ "tf.reduce_mean(phi_xs * score + grad_phi_x).numpy() ** 2" ], "execution_count": 8, "outputs": [ { "output_type": "execute_result", "data": { "text/plain": [ "5.552772948637223e-05" ] }, "metadata": { "tags": [] }, "execution_count": 8 } ] }, @@ -291,6 +291,17 @@ "source": [ "Indeed approaches 0!" ] }, { "cell_type": "code", "metadata": { "id": "c7X61SRtvrCQ" }, "source": [ "" ], "execution_count": 8, "outputs": [] } ] } -
Georges Pantalos revised this gist
Feb 18, 2021 . 1 changed file with 9 additions and 9 deletions.There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode charactersOriginal file line number Diff line number Diff line change @@ -6,7 +6,7 @@ "name": "stein_identity.ipynb", "provenance": [], "collapsed_sections": [], "authorship_tag": "ABX9TyNB8oQc1T3c7sxMsXPW2CLH", "include_colab_link": true }, "kernelspec": { @@ -50,7 +50,7 @@ "tfd = tfp.distributions\n", "tf.random.set_seed(0)" ], "execution_count": null, "outputs": [] }, { @@ -59,7 +59,7 @@ "id": "dny1VJB9V3nL" }, "source": [ "Define a distribution $p = \\mathcal N(0,1)$" ] }, { @@ -75,7 +75,7 @@ "p = tfd.Normal(0,1)\n", "p" ], "execution_count": null, "outputs": [ { "output_type": "execute_result", @@ -114,7 +114,7 @@ "x = tf.Variable(x)\n", "x.shape" ], "execution_count": null, "outputs": [ { "output_type": "execute_result", @@ -154,7 +154,7 @@ "score = tape.gradient(logp, x)\n", "score.shape" ], "execution_count": null, "outputs": [ { "output_type": "execute_result", @@ -194,7 +194,7 @@ " # add smooth functions...\n", "]" ], "execution_count": null, "outputs": [] }, { @@ -230,7 +230,7 @@ "grad_phi_x = tf.stack(grad_phi_x)\n", "grad_phi_x.shape" ], "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": null, "outputs": [ { "output_type": "execute_result", -
Georges Pantalos revised this gist
Feb 18, 2021 . No changes.There are no files selected for viewing
-
Georges Pantalos revised this gist
Feb 18, 2021 . 1 changed file with 12 additions and 1 deletion.There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode charactersOriginal 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", "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": { -
Georges Pantalos created this gist
Feb 18, 2021 .There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode charactersOriginal 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!" ] } ] }