Skip to content

Instantly share code, notes, and snippets.

@Sayam753
Created July 26, 2020 19:39
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save Sayam753/f434492fc19f78bb93f3002cdecfd002 to your computer and use it in GitHub Desktop.
distributions.ipynb
Display the source blob
Display the rendered blob
Raw
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.7"
},
"colab": {
"name": "distributions.ipynb",
"provenance": [],
"collapsed_sections": [],
"include_colab_link": true
}
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/gist/Sayam753/f434492fc19f78bb93f3002cdecfd002/distributions.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "jhIEqOu_jw7K",
"colab_type": "text"
},
"source": [
"# Distributions Enhancement Proposal - PyMC4\n",
"\n",
"## Add transformations to bounded distributions\n",
"\n",
"At the current state of PyMC4, only `Log` and `Logit` transformations are included to do `mcmc` sampling or `vi` approximations. No transformations have been there for Bounded distributions. This notebook explores some ways of adding `Interval` transform."
]
},
{
"cell_type": "code",
"metadata": {
"id": "8dFyetCwjw7M",
"colab_type": "code",
"colab": {},
"outputId": "8528de2b-45dd-4699-e3a8-2f08fa505a2f"
},
"source": [
"!git checkout master"
],
"execution_count": null,
"outputs": [
{
"output_type": "stream",
"text": [
"Switched to branch 'master'\n",
"Your branch is up to date with 'origin/master'.\n"
],
"name": "stdout"
}
]
},
{
"cell_type": "code",
"metadata": {
"id": "553A1jEijw7S",
"colab_type": "code",
"colab": {}
},
"source": [
"import arviz as az\n",
"import math\n",
"import numpy as np\n",
"import pymc4\n",
"import tensorflow as tf\n",
"import tensorflow_probability as tfp\n",
"\n",
"tfd = tfp.distributions\n",
"tfb = tfp.bijectors"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "c2IslEDajw7W",
"colab_type": "code",
"colab": {}
},
"source": [
"experiments = np.array([1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.])\n",
"\n",
"@pymc4.model\n",
"def model():\n",
" prob = yield pymc4.Uniform('p', 0., 1.)\n",
" ll = yield pymc4.Bernoulli('ll', prob, observed=experiments)"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "sHb430hwjw7Z",
"colab_type": "code",
"colab": {},
"outputId": "c8275da6-cec3-487f-d238-f60362ed2b76"
},
"source": [
"advi = pymc4.fit(model())"
],
"execution_count": null,
"outputs": [
{
"output_type": "stream",
"text": [
"WARNING:tensorflow:From /usr/local/lib/python3.7/site-packages/tensorflow_probability/python/math/minimize.py:77: calling <lambda> (from tensorflow_probability.python.vi.optimization) with loss is deprecated and will be removed after 2020-07-01.\n",
"Instructions for updating:\n",
"The signature for `trace_fn`s passed to `minimize` has changed. Trace functions now take a single `traceable_quantities` argument, which is a `tfp.math.MinimizeTraceableQuantities` namedtuple containing `traceable_quantities.loss`, `traceable_quantities.gradients`, etc. Please update your `trace_fn` definition.\n"
],
"name": "stdout"
}
]
},
{
"cell_type": "code",
"metadata": {
"id": "uRmPtUI2jw7c",
"colab_type": "code",
"colab": {},
"outputId": "63c88dc3-e9bf-40f2-c921-088c1892019c"
},
"source": [
"np.sum(np.isnan(advi.losses)) # there seems are a large section of ELBO loss containing nans"
],
"execution_count": null,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"6867"
]
},
"metadata": {
"tags": []
},
"execution_count": 5
}
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "c-lpS17Kjw7h",
"colab_type": "text"
},
"source": [
"One can also pass `validate_args=True` to each distribution while doing VI, and justify support matching constraint failing.\n",
"\n",
"## Possible Solutions\n",
"## 1. Using conditionals"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "vliJCRLtjw7i",
"colab_type": "text"
},
"source": [
"### Transformations are defined as -\n",
"\n",
"```python\n",
"from pymc4.distributions.transforms import JacobianPreference, Transform\n",
"\n",
"class Interval(Transform):\n",
" name = \"interval\"\n",
" JacobianPreference = JacobianPreference.Backward\n",
"\n",
" def __init__(self, lower_limit, upper_limit):\n",
" if math.isinf(lower_limit) and math.isinf(upper_limit):\n",
" transform = tfb.Identity()\n",
" \n",
" elif math.isinf(lower_limit) and math.isfinite(upper_limit):\n",
" transform = tfb.Chain([tfb.Shift(upper_limit), tfb.Scale(-1), tfb.Exp()]) # upper - exp(x)\n",
" \n",
" elif math.isfinite(lower_limit) and math.isinf(upper_limit):\n",
" transform = tfb.Chain([tfb.Shift(lower_limit), tfb.Exp()]) # exp(x) + lower\n",
" \n",
" else:\n",
" transform = tfb.Sigmoid(low=lower_limit, high=upper_limit) # interval\n",
" self._transform = transform\n",
"\n",
" def forward(self, x):\n",
" return self._transform.inverse(x)\n",
"\n",
" def inverse(self, z):\n",
" return self._transform.forward(z)\n",
"\n",
" def forward_log_det_jacobian(self, x):\n",
" return self._transform.inverse_log_det_jacobian(x, self._transform.inverse_min_event_ndims)\n",
"\n",
" def inverse_log_det_jacobian(self, z):\n",
" return self._transform.forward_log_det_jacobian(z, self._transform.forward_min_event_ndims)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8cW8QiArjw7i",
"colab_type": "text"
},
"source": [
"### Fitting the model -\n",
"\n",
"```python\n",
"advi = pymc4.fit(model())\n",
"\n",
"# Traceback\n",
"~/Desktop/pymc/pymc4/pymc4/distributions/distribution.py in _init_transform(self, transform)\n",
" 280 def _init_transform(self, transform):\n",
" 281 if transform is None:\n",
"--> 282 return transforms.Interval(self.lower_limit(), self.upper_limit())\n",
" 283 else:\n",
" 284 return transform\n",
"\n",
"~/Desktop/pymc/pymc4/pymc4/distributions/transforms.py in __init__(self, lower_limit, upper_limit)\n",
" 160 \n",
" 161 def __init__(self, lower_limit, upper_limit):\n",
"--> 162 if math.isinf(lower_limit) and math.isinf(upper_limit):\n",
" 163 transform = tfb.Identity()\n",
" 164 \n",
"\n",
"TypeError: must be real number, not Tensor\n",
"\n",
"```\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "kjHvEnzGjw7j",
"colab_type": "text"
},
"source": [
"## 2. Using tf.cond"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "HIsy-213jw7k",
"colab_type": "text"
},
"source": [
"### Transformations are defined as -\n",
"\n",
"```python\n",
"from pymc4.distributions.transforms import JacobianPreference, Transform\n",
"\n",
"class Interval(Transform):\n",
" name = \"interval\"\n",
" JacobianPreference = JacobianPreference.Backward\n",
"\n",
" def __init__(self, lower_limit, upper_limit):\n",
" transform = tf.cond(\n",
" tf.math.is_inf(lower_limit),\n",
" lambda: tf.cond(\n",
" tf.math.is_inf(upper_limit),\n",
" lambda: tfb.Identity(),\n",
" lambda: tfb.Chain([tfb.Shift(upper_limit), tfb.Scale(-1), tfb.Exp()]), # upper - exp(x)\n",
" ),\n",
" lambda: tf.cond(\n",
" tf.math.is_inf(upper_limit),\n",
" lambda: tfb.Chain([tfb.Shift(lower_limit), tfb.Exp()]), # exp(x) + lower\n",
" lambda: tfb.Sigmoid(low=lower_limit, high=upper_limit), # interval\n",
" ),\n",
" )\n",
" self._transform = transform\n",
"\n",
" def forward(self, x):\n",
" return self._transform.inverse(x)\n",
"\n",
" def inverse(self, z):\n",
" return self._transform.forward(z)\n",
"\n",
" def forward_log_det_jacobian(self, x):\n",
" return self._transform.inverse_log_det_jacobian(x, self._transform.inverse_min_event_ndims)\n",
"\n",
" def inverse_log_det_jacobian(self, z):\n",
" return self._transform.forward_log_det_jacobian(z, self._transform.forward_min_event_ndims)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "xXmtXsJWjw7k",
"colab_type": "text"
},
"source": [
"### Fitting the model -\n",
"\n",
"```python\n",
"advi = pymc4.fit(model())\n",
"\n",
"# Traceback\n",
"/usr/local/lib/python3.7/site-packages/tensorflow/python/util/nest.py in <listcomp>(.0)\n",
" 635 \n",
" 636 return pack_sequence_as(\n",
"--> 637 structure[0], [func(*x) for x in entries],\n",
" 638 expand_composites=expand_composites)\n",
" 639 \n",
"\n",
"/usr/local/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py in convert(x)\n",
" 946 \"must return zero or more Tensors; in compilation of %s, found \"\n",
" 947 \"return value of type %s, which is not a Tensor.\" %\n",
"--> 948 (str(python_func), type(x)))\n",
" 949 if add_control_dependencies:\n",
" 950 x = deps_ctx.mark_as_return(x)\n",
"\n",
"TypeError: To be compatible with tf.eager.defun, Python functions must return zero or more Tensors; in compilation of <function Interval.__init__.<locals>.<lambda>.<locals>.<lambda> at 0x137ceeef0>, found return value of type <class 'tensorflow_probability.python.bijectors.identity.Identity'>, which is not a Tensor.\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "tKXxbj64jw7l",
"colab_type": "text"
},
"source": [
"## 3. Manually inspect the distributions and add transformations\n",
"This manual addition has already been included in the PR [#289](https://github.com/pymc-devs/pymc4/pull/289) Add Full Rank Approximation to PyMC4. ADVI works but this process does not seem very natural."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Lj4LPgd_jw7l",
"colab_type": "text"
},
"source": [
"## 4. Proposed Approach -\n",
"It is difficult to propose a generic solution as we do not know which values to transform beforehand. To solve this, classify the Bounded distributions as -\n",
"1. Upper Bounded Distributions\n",
"2. Lower Bounded Distributions\n",
"3. Interval Bounded Distributions\n",
"\n",
"After adding transformations to these base classes, inherit them to their corresponding distributions. This approach has been proposed after discussions with Maxim.\n",
"\n",
"This proposed implementation is ready to test on a separate branch [https://github.com/Sayam753/pymc4/tree/proposed_transformations](https://github.com/Sayam753/pymc4/tree/proposed_transformations). Any feedback over this.\n",
"\n",
"Also, this discussion in google groups (https://groups.google.com/a/tensorflow.org/forum/#!topic/tfprobability/us6ZzR_WTZU) mentions about automatically applying bijectors in future from TFP side."
]
}
]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment