Merge pull request #2416 from ferrine/approximations_refactoring

Refactoring to support Structured VI

Author: twiecki

docs/source/notebooks/GLM-hierarchical-advi-minibatch.ipynb

@@ -31,7 +31,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "%autosave 0\n",
@@ -121,7 +123,9 @@
   {
    "cell_type": "code",
    "execution_count": 4,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "from keras.models import Sequential\n",
@@ -172,7 +176,9 @@
   {
    "cell_type": "code",
    "execution_count": 8,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "def cnn_enc(xs, latent_dim, nb_filters=64, nb_conv=3, intermediate_dim=128):\n",
@@ -216,7 +222,9 @@
   {
    "cell_type": "code",
    "execution_count": 9,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "class Encoder:\n",
@@ -292,7 +300,9 @@
   {
    "cell_type": "code",
    "execution_count": 10,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "def cnn_dec(zs, nb_filters=64, nb_conv=3, output_shape=(1, 28, 28)):\n",
@@ -331,7 +341,9 @@
   {
    "cell_type": "code",
    "execution_count": 11,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "class Decoder:\n",
@@ -385,7 +397,9 @@
   {
    "cell_type": "code",
    "execution_count": 12,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "# Constants\n",
@@ -454,7 +468,9 @@
   {
    "cell_type": "code",
    "execution_count": 14,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "with pm.Model() as model:\n",
@@ -478,10 +494,12 @@
   {
    "cell_type": "code",
    "execution_count": 15,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
-    "local_RVs = OrderedDict({zs: (enc.means, enc.rhos)})"
+    "local_RVs = OrderedDict({zs: dict(mu=enc.means, rho=enc.rhos)})"
    ]
   },
   {
@@ -671,7 +689,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.1"
+   "version": "3.6.0b4"
   },
   "latex_envs": {
    "bibliofile": "biblio.bib",

docs/source/notebooks/api_quickstart.ipynb

@@ -4,10 +4,12 @@
 creating custom backends).
 """
 import numpy as np
-from ..model import modelcontext
 import warnings
 import theano.tensor as tt
 
+from ..model import modelcontext
+
+
 class BackendError(Exception):
     pass
 
@@ -24,11 +26,13 @@ class BaseTrace(object):
     vars : list of variables
         Sampling values will be stored for these variables. If None,
         `model.unobserved_RVs` is used.
+    test_point : dict
+        use different test point that might be with changed variables shapes
     """
 
     supports_sampler_stats = False
 
-    def __init__(self, name, model=None, vars=None):
+    def __init__(self, name, model=None, vars=None, test_point=None):
         self.name = name
 
         model = modelcontext(model)
@@ -41,7 +45,13 @@ def __init__(self, name, model=None, vars=None):
 
         # Get variable shapes. Most backends will need this
         # information.
-        var_values = list(zip(self.varnames, self.fn(model.test_point)))
+        if test_point is None:
+            test_point = model.test_point
+        else:
+            test_point_ = model.test_point.copy()
+            test_point_.update(test_point)
+            test_point = test_point_
+        var_values = list(zip(self.varnames, self.fn(test_point)))
         self.var_shapes = {var: value.shape
                            for var, value in var_values}
         self.var_dtypes = {var: value.dtype
@@ -72,7 +82,6 @@ def _set_sampler_vars(self, sampler_vars):
 
         self.sampler_vars = sampler_vars
 
-
     def setup(self, draws, chain, sampler_vars=None):
         """Perform chain-specific setup.
 

docs/source/notebooks/bayesian_neural_network_advi.ipynb

@@ -18,24 +18,26 @@ def activator(instance):
 class HDF5(base.BaseTrace):
     """HDF5 trace object
 
-        Parameters
-        ----------
-        name : str
-            Name of backend. This has no meaning for the HDF5 backend.
-        model : Model
-            If None, the model is taken from the `with` context.
-        vars : list of variables
-            Sampling values will be stored for these variables. If None,
-            `model.unobserved_RVs` is used.
+    Parameters
+    ----------
+    name : str
+        Name of backend. This has no meaning for the HDF5 backend.
+    model : Model
+        If None, the model is taken from the `with` context.
+    vars : list of variables
+        Sampling values will be stored for these variables. If None,
+        `model.unobserved_RVs` is used.
+    test_point : dict
+        use different test point that might be with changed variables shapes
         """
 
     supports_sampler_stats = True
 
-    def __init__(self, name=None, model=None, vars=None):
+    def __init__(self, name=None, model=None, vars=None, test_point=None):
         self.hdf5_file = None
         self.draw_idx = 0
         self.draws = None
-        super(HDF5, self).__init__(name, model, vars)
+        super(HDF5, self).__init__(name, model, vars, test_point)
 
     def _get_sampler_stats(self, varname, sampler_idx, burn, thin):
         with self.activate_file:

docs/source/notebooks/convolutional_vae_keras_advi.ipynb

@@ -22,8 +22,8 @@ class NDArray(base.BaseTrace):
 
     supports_sampler_stats = True
 
-    def __init__(self, name=None, model=None, vars=None):
-        super(NDArray, self).__init__(name, model, vars)
+    def __init__(self, name=None, model=None, vars=None, test_point=None):
+        super(NDArray, self).__init__(name, model, vars, test_point)
         self.draw_idx = 0
         self.draws = None
         self.samples = {}

docs/source/notebooks/empirical-approx-overview.ipynb

@@ -70,10 +70,12 @@ class SQLite(base.BaseTrace):
     vars : list of variables
         Sampling values will be stored for these variables. If None,
         `model.unobserved_RVs` is used.
+    test_point : dict
+        use different test point that might be with changed variables shapes
     """
 
-    def __init__(self, name, model=None, vars=None):
-        super(SQLite, self).__init__(name, model, vars)
+    def __init__(self, name, model=None, vars=None, test_point=None):
+        super(SQLite, self).__init__(name, model, vars, test_point)
         self._var_cols = {}
         self.var_inserts = {}  # varname -> insert statement
         self.draw_idx = 0

docs/source/notebooks/gaussian-mixture-model-advi.ipynb

@@ -36,12 +36,14 @@ class Text(base.BaseTrace):
     vars : list of variables
         Sampling values will be stored for these variables. If None,
         `model.unobserved_RVs` is used.
+    test_point : dict
+        use different test point that might be with changed variables shapes
     """
 
-    def __init__(self, name, model=None, vars=None):
+    def __init__(self, name, model=None, vars=None, test_point=None):
         if not os.path.exists(name):
             os.mkdir(name)
-        super(Text, self).__init__(name, model, vars)
+        super(Text, self).__init__(name, model, vars, test_point)
 
         self.flat_names = {v: ttab.create_flat_names(v, shape)
                            for v, shape in self.var_shapes.items()}

docs/source/notebooks/gaussian_process.ipynb

@@ -23,28 +23,26 @@ class ArrayOrdering(object):
 
     def __init__(self, vars):
         self.vmap = []
-        self._by_name = {}
-        size = 0
+        self.by_name = {}
+        self.size = 0
 
         for var in vars:
             name = var.name
             if name is None:
                 raise ValueError('Unnamed variable in ArrayOrdering.')
-            if name in self._by_name:
+            if name in self.by_name:
                 raise ValueError('Name of variable not unique: %s.' % name)
             if not hasattr(var, 'dshape') or not hasattr(var, 'dsize'):
                 raise ValueError('Shape of variable not known %s' % name)
 
-            slc = slice(size, size + var.dsize)
+            slc = slice(self.size, self.size + var.dsize)
             varmap = VarMap(name, slc, var.dshape, var.dtype)
             self.vmap.append(varmap)
-            self._by_name[name] = varmap
-            size += var.dsize
-
-        self.size = size
+            self.by_name[name] = varmap
+            self.size += var.dsize
 
     def __getitem__(self, key):
-        return self._by_name[key]
+        return self.by_name[key]
 
 
 class DictToArrayBijection(object):
@@ -122,24 +120,22 @@ class ListArrayOrdering(object):
     """
 
     def __init__(self, list_arrays, intype='numpy'):
+        if intype not in {'tensor', 'numpy'}:
+            raise ValueError("intype not in {'tensor', 'numpy'}")
         self.vmap = []
-        dim = 0
-
-        count = 0
+        self.intype = intype
+        self.size = 0
         for array in list_arrays:
-            if intype == 'tensor':
+            if self.intype == 'tensor':
                 name = array.name
                 array = array.tag.test_value
-            elif intype == 'numpy':
+            else:
                 name = 'numpy'
 
-            slc = slice(dim, dim + array.size)
+            slc = slice(self.size, self.size + array.size)
             self.vmap.append(DataMap(
-                count, slc, array.shape, array.dtype, name))
-            dim += array.size
-            count += 1
-
-        self.size = dim
+                len(self.vmap), slc, array.shape, array.dtype, name))
+            self.size += array.size
 
 
 class ListToArrayBijection(object):

docs/source/notebooks/lda-advi-aevb.ipynb

@@ -91,7 +91,7 @@ def i0(x):
     return tt.switch(tt.lt(x, 5), 1. + x**2 / 4. + x**4 / 64. + x**6 / 2304. + x**8 / 147456.
                      + x**10 / 14745600. + x**12 / 2123366400.,
                      np.e**x / (2. * np.pi * x)**0.5 * (1. + 1. / (8. * x) + 9. / (128. * x**2) + 225. / (3072 * x**3)
-                                                       + 11025. / (98304. * x**4)))
+                                                        + 11025. / (98304. * x**4)))
 
 
 def i1(x):
@@ -108,14 +108,14 @@ def sd2rho(sd):
     """
     `sd -> rho` theano converter
     :math:`mu + sd*e = mu + log(1+exp(rho))*e`"""
-    return tt.log(tt.exp(sd) - 1.)
+    return tt.log(tt.exp(tt.abs_(sd)) - 1.)
 
 
 def rho2sd(rho):
     """
     `rho -> sd` theano converter
     :math:`mu + sd*e = mu + log(1+exp(rho))*e`"""
-    return tt.log1p(tt.exp(rho))
+    return tt.nnet.softplus(rho)
 
 
 def log_normal(x, mean, **kwargs):