From c94f6ff75bb5f9b2894c999cd40c35b8c5d89583 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Caba=C3=B1as?= <rcabanas@idsia.ch>
Date: Sat, 11 Jul 2020 20:36:48 +0200
Subject: [PATCH 01/11] PGM example

---
 examples/probzoo/nl_digit_classifier.py | 112 ++++++++++++++++++++++++
 1 file changed, 112 insertions(+)
 create mode 100644 examples/probzoo/nl_digit_classifier.py

diff --git a/examples/probzoo/nl_digit_classifier.py b/examples/probzoo/nl_digit_classifier.py
new file mode 100644
index 0000000..aa663cd
--- /dev/null
+++ b/examples/probzoo/nl_digit_classifier.py
@@ -0,0 +1,112 @@
+# required pacakges
+import inferpy as inf
+import numpy as np
+import tensorflow as tf
+
+import matplotlib.pyplot as plt
+from inferpy.data import mnist
+
+
+# number of observations (dataset size)
+N = 1000
+
+# batch size
+M = 100
+
+
+scale_epsilon = 0.01
+
+
+@inf.probmodel
+def digit_classifier(k, d0, dx, dy):
+
+    with inf.datamodel():
+        z = inf.Normal(tf.ones(k) * 0.1, 1., name="z")
+
+        decoder = tf.keras.Sequential([
+            tf.keras.layers.Dense(d0, tf.nn.relu),
+            tf.keras.layers.Dense(dx)
+        ])
+
+        classifier = tf.keras.Sequential([
+            tf.keras.layers.Dense(dy)
+        ])
+
+        x = inf.Normal(decoder(z), 1., name="x")
+        y = inf.Categorical(logits=classifier(z), name="y")
+
+p = digit_classifier(k=2, d0=100, dx=28*28, dy=3)
+
+
+@inf.probmodel
+def qmodel(k, d0, dx):
+    with inf.datamodel():
+        x = inf.Normal(tf.ones(dx), 1, name="x")
+
+        encoder = tf.keras.Sequential([
+            tf.keras.layers.Dense(d0, activation=tf.nn.relu),
+            tf.keras.layers.Dense(2 * k)
+        ])
+
+        output = encoder(x)
+        qz_loc = output[:, :k]
+        qz_scale = tf.nn.softplus(output[:, k:])+0.01
+        qz = inf.Normal(qz_loc, qz_scale, name="z")
+
+q = qmodel(k=2, d0=100, dx=28*28)
+
+# get the MNIST dataset
+(x_train, y_train), (x_test, y_test) = mnist.load_data(num_instances=N, digits=[0, 1, 2])
+
+# set the inference algorithm
+SVI = inf.inference.SVI(q, epochs=5000, batch_size=M)
+
+# fit the model to the data
+p.fit({"x": x_train, "y":y_train}, SVI)
+
+
+# Plot the evolution of the loss
+L = SVI.losses
+plt.plot(range(len(L)), L)
+
+plt.xlabel('epochs')
+plt.ylabel('Loss')
+plt.title('Loss evolution')
+plt.grid(True)
+plt.show()
+
+
+# Extract the posterior of z given the training data
+postz = np.concatenate([
+    p.posterior("z", data={"x": x_train[i:i + M, :]}).sample()
+    for i in range(0,N,M)])
+
+# for each input instance, plot the hidden encoding coloured by the number that it represents
+markers = ["x", "+", "o"]
+colors = [plt.get_cmap("gist_rainbow")(0.05),
+          plt.get_cmap("gnuplot2")(0.08),
+          plt.get_cmap("gist_rainbow")(0.33)]
+transp = [0.9, 0.9, 0.5]
+
+fig = plt.figure()
+
+for c in range(0, len(DIG)):
+    col = colors[c]
+    plt.scatter(postz[y_train == DIG[c], 0], postz[y_train == DIG[c], 1], color=col,
+                label=DIG[c], marker=markers[c], alpha=transp[c], s=60)
+    plt.legend()
+
+plt.show()
+
+
+
+# Predict a set of images
+def predict(x):
+    postz = p.posterior("z", data={"x": x}).sample()
+    return p.posterior_predictive("y", data={"z":postz}).sample()
+
+y_gen = predict(x_test[:M])
+
+# Compute the accuracy
+acc = np.sum(y_test[:M] == y_gen)/M
+print(f"accuracy: {acc}")
\ No newline at end of file

From e1bdcd3425cb4caabde217ccdc7d992a2ad94c32 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Caba=C3=B1as?= <rcabanas@idsia.ch>
Date: Wed, 30 Sep 2020 15:32:10 +0200
Subject: [PATCH 02/11] [FEAT] test data of different size

---
 examples/probzoo/nl_digit_classifier.py |  20 ++---
 inferpy/queries/query.py                | 102 +++++++++++++++++++-----
 2 files changed, 89 insertions(+), 33 deletions(-)

diff --git a/examples/probzoo/nl_digit_classifier.py b/examples/probzoo/nl_digit_classifier.py
index aa663cd..97478a9 100644
--- a/examples/probzoo/nl_digit_classifier.py
+++ b/examples/probzoo/nl_digit_classifier.py
@@ -14,12 +14,8 @@
 M = 100
 
 
-scale_epsilon = 0.01
-
-
 @inf.probmodel
 def digit_classifier(k, d0, dx, dy):
-
     with inf.datamodel():
         z = inf.Normal(tf.ones(k) * 0.1, 1., name="z")
 
@@ -27,11 +23,9 @@ def digit_classifier(k, d0, dx, dy):
             tf.keras.layers.Dense(d0, tf.nn.relu),
             tf.keras.layers.Dense(dx)
         ])
-
         classifier = tf.keras.Sequential([
             tf.keras.layers.Dense(dy)
         ])
-
         x = inf.Normal(decoder(z), 1., name="x")
         y = inf.Categorical(logits=classifier(z), name="y")
 
@@ -47,7 +41,6 @@ def qmodel(k, d0, dx):
             tf.keras.layers.Dense(d0, activation=tf.nn.relu),
             tf.keras.layers.Dense(2 * k)
         ])
-
         output = encoder(x)
         qz_loc = output[:, :k]
         qz_scale = tf.nn.softplus(output[:, k:])+0.01
@@ -56,11 +49,10 @@ def qmodel(k, d0, dx):
 q = qmodel(k=2, d0=100, dx=28*28)
 
 # get the MNIST dataset
-(x_train, y_train), (x_test, y_test) = mnist.load_data(num_instances=N, digits=[0, 1, 2])
-
+(x_train, y_train), (x_test, y_test) = mnist.load_data(
+    num_instances=N, digits=[0, 1, 2])
 # set the inference algorithm
-SVI = inf.inference.SVI(q, epochs=5000, batch_size=M)
-
+SVI = inf.inference.SVI(q, epochs=10000, batch_size=M)
 # fit the model to the data
 p.fit({"x": x_train, "y":y_train}, SVI)
 
@@ -76,7 +68,7 @@ def qmodel(k, d0, dx):
 plt.show()
 
 
-# Extract the posterior of z given the training data
+# extract the posterior of z given the training data
 postz = np.concatenate([
     p.posterior("z", data={"x": x_train[i:i + M, :]}).sample()
     for i in range(0,N,M)])
@@ -100,13 +92,13 @@ def qmodel(k, d0, dx):
 
 
 
-# Predict a set of images
+# predict a set of images
 def predict(x):
     postz = p.posterior("z", data={"x": x}).sample()
     return p.posterior_predictive("y", data={"z":postz}).sample()
 
 y_gen = predict(x_test[:M])
 
-# Compute the accuracy
+# compute the accuracy
 acc = np.sum(y_test[:M] == y_gen)/M
 print(f"accuracy: {acc}")
\ No newline at end of file
diff --git a/inferpy/queries/query.py b/inferpy/queries/query.py
index 8697e97..8d64c3f 100644
--- a/inferpy/queries/query.py
+++ b/inferpy/queries/query.py
@@ -3,6 +3,8 @@
 
 from inferpy import contextmanager
 from inferpy import util
+import math
+import tensorflow as tf
 
 
 def flatten_result(f):
@@ -35,15 +37,51 @@ def __init__(self, variables, target_names=None, data={}, enable_interceptor_var
         self.data = data
         self.enable_interceptor_variables = enable_interceptor_variables
 
+
+        self.batch_size = [v for (k,v) in self.observed_variables.items() if v.is_datamodel][0].shape[0].value
+        self.data_size = self.data[[k for (k,v) in self.observed_variables.items() if k in self.data.keys() and v.is_datamodel][0]].shape[0]
+
+        ####
+
+
+        rows_to_add = round((math.ceil(self.data_size / self.batch_size) - self.data_size / self.batch_size) * self.batch_size)
+        num_batches = math.ceil(self.data_size / self.batch_size)
+
+        for k,v in self.observed_variables.items():
+            if k in self.data and v.is_datamodel:
+                padding = np.zeros([len(data[k].shape), 2], np.int32)
+                padding[0, 1] = rows_to_add
+                data[k] = np.pad(data[k], padding, mode="constant")
+
+        self.batches = [{} for _ in range(num_batches)]
+
+        k,v = list(self.observed_variables.items())[0]
+        for k,v in self.observed_variables.items():
+            if k in self.data and v.is_datamodel:
+                batches_k = np.split(data[k], num_batches, axis=0)
+                for i in range(num_batches):
+                    self.batches[i][k] = batches_k[i]
+            elif k in self.data:
+                for i in range(num_batches):
+                    self.batches[i][k] = data[k]
+
+
+
     @flatten_result
     @util.tf_run_ignored
     def log_prob(self):
         """ Computes the log probabilities of a (set of) sample(s)"""
-        with util.interceptor.enable_interceptor(*self.enable_interceptor_variables):
-            with contextmanager.observe(self.observed_variables, self.data):
-                result = util.runtime.try_run({k: v.log_prob(v.value) for k, v in self.target_variables.items()})
 
-        return result
+        results = []
+        for batch in self.batches:
+            with util.interceptor.enable_interceptor(*self.enable_interceptor_variables):
+                with contextmanager.observe(self.observed_variables, batch):
+                    result = util.runtime.try_run({k: v.log_prob(v.value) for k, v in self.target_variables.items()})
+
+            results.append(result)
+
+
+        return {k:np.vstack([r[k] for r in results])[:self.data_size] for k in self.target_variables.keys()}
 
     def sum_log_prob(self):
         """ Computes the sum of the log probabilities (evaluated) of a (set of) sample(s)"""
@@ -54,19 +92,26 @@ def sum_log_prob(self):
     @util.tf_run_ignored
     def sample(self, size=1):
         """ Generates a sample for eache variable in the model """
-        with util.interceptor.enable_interceptor(*self.enable_interceptor_variables):
-            with contextmanager.observe(self.observed_variables, self.data):
-                # each iteration for `size` run the dict in the session, so if there are dependencies among random vars
-                # they are computed in the same graph operations, and reflected in the results
-                samples = [util.runtime.try_run(self.target_variables) for _ in range(size)]
-
-        if size == 1:
-            result = samples[0]
-        else:
-            # compact all samples in one single dict
-            result = {k: np.array([sample[k] for sample in samples]) for k in self.target_variables.keys()}
 
-        return result
+        results = []
+        for batch in self.batches:
+            with util.interceptor.enable_interceptor(*self.enable_interceptor_variables):
+                with contextmanager.observe(self.observed_variables, batch):
+                    # each iteration for `size` run the dict in the session, so if there are dependencies among random vars
+                    # they are computed in the same graph operations, and reflected in the results
+                    samples = [util.runtime.try_run(self.target_variables) for _ in range(size)]
+
+            if size == 1:
+                result = samples[0]
+            else:
+                # compact all samples in one single dict
+                result = {k: np.array([sample[k] for sample in samples]) for k in self.target_variables.keys()}
+
+            results.append(result)
+
+
+        return {k:np.vstack([r[k] for r in results])[:self.data_size] for k in self.target_variables.keys()}
+
 
     @flatten_result
     @util.tf_run_ignored
@@ -104,8 +149,27 @@ def filter_parameters(varname, parameters):
 
             return {k: util.runtime.try_run(v) for k, v in parameters.items() if k in selected_parameters}
 
-        with contextmanager.observe(self.observed_variables, self.data):
-            result = {k: filter_parameters(k, v.parameters)
-                      for k, v in self.target_variables.items()}
+        # function that merges the parameters of 2 batches
+        def merge_params(p1, p2):
+            out = {}
+            for k,v in p1.items():
+                if np.ndim(v) == 0:
+                    out[k] = v
+                else:
+                    out[k] = np.vstack([p1[k], p2[k]])[0:self.data_size]
+            return out
+
+        # get the parameter for each batch
+        result = {}
+
+        for batch in self.batches:
+            with contextmanager.observe(self.observed_variables, batch):
+                r = {k: filter_parameters(k, v.parameters)
+                          for k, v in self.target_variables.items()}
+
+            if len(result.keys()) == 0:
+                result = r
+            else:
+                result = {var:merge_params(result[var], r[var]) for var in self.target_variables.keys()}
 
         return result

From c115095ea48889c93d8bb9791d5e09b76da8f294 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Caba=C3=B1as?= <rcabanas@idsia.ch>
Date: Thu, 15 Oct 2020 10:14:52 +0200
Subject: [PATCH 03/11] test_queries

---
 examples/docs/guideinference/1.py |   6 +-
 inferpy/data/mnist.py             |  12 +-
 inferpy/queries/query.py          | 102 ++++-------------
 inferpy/util/common.py            |  30 +++--
 inferpy/util/iterables.py         |  15 +--
 tests/test_queries.py             | 178 ++++++++++++++++++++++++++++++
 6 files changed, 221 insertions(+), 122 deletions(-)
 create mode 100644 tests/test_queries.py

diff --git a/examples/docs/guideinference/1.py b/examples/docs/guideinference/1.py
index f9ac740..2345b97 100644
--- a/examples/docs/guideinference/1.py
+++ b/examples/docs/guideinference/1.py
@@ -104,8 +104,4 @@ def custom_elbo(pvars, qvars, **kwargs):
 #104
 
 # extract the posterior of z
-hidden_encoding = m.posterior("z").parameters()["samples"].mean(axis=0)
-
-
-
-SVI = inf.inference.SVI(qmodel(k=1,d=2), epochs=1000, batch_size=200)
\ No newline at end of file
+hidden_encoding = m.posterior("z").parameters()["samples"].mean(axis=0)
\ No newline at end of file
diff --git a/inferpy/data/mnist.py b/inferpy/data/mnist.py
index 045b21e..b40f33a 100644
--- a/inferpy/data/mnist.py
+++ b/inferpy/data/mnist.py
@@ -7,14 +7,10 @@
 def load_data(vectorize=True, num_instances=None, num_instances_test=None, digits=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]):
     """ Loads the MNIST datase
 
-        Args:
-            vectorize: if true, each 2D image is transformed into a 1D vector
-            num_instances: total number of images loaded
-            digits: list of integers indicating the digits to be considered
-
-        Returns:
-            Tuple of Numpy arrays: `(x_train, y_train), (x_test, y_test)
-
+    :param vectorize: if true, each 2D image is transformed into a 1D vector
+    :param num_instances: total number of images loaded
+    :param digits: list of integers indicating the digits to be considered
+    :return:  Tuple of Numpy arrays: `(x_train, y_train), (x_test, y_test)
     """
 
     (x_train, y_train), (x_test, y_test) = mnist.load_data()
diff --git a/inferpy/queries/query.py b/inferpy/queries/query.py
index 8d64c3f..8697e97 100644
--- a/inferpy/queries/query.py
+++ b/inferpy/queries/query.py
@@ -3,8 +3,6 @@
 
 from inferpy import contextmanager
 from inferpy import util
-import math
-import tensorflow as tf
 
 
 def flatten_result(f):
@@ -37,51 +35,15 @@ def __init__(self, variables, target_names=None, data={}, enable_interceptor_var
         self.data = data
         self.enable_interceptor_variables = enable_interceptor_variables
 
-
-        self.batch_size = [v for (k,v) in self.observed_variables.items() if v.is_datamodel][0].shape[0].value
-        self.data_size = self.data[[k for (k,v) in self.observed_variables.items() if k in self.data.keys() and v.is_datamodel][0]].shape[0]
-
-        ####
-
-
-        rows_to_add = round((math.ceil(self.data_size / self.batch_size) - self.data_size / self.batch_size) * self.batch_size)
-        num_batches = math.ceil(self.data_size / self.batch_size)
-
-        for k,v in self.observed_variables.items():
-            if k in self.data and v.is_datamodel:
-                padding = np.zeros([len(data[k].shape), 2], np.int32)
-                padding[0, 1] = rows_to_add
-                data[k] = np.pad(data[k], padding, mode="constant")
-
-        self.batches = [{} for _ in range(num_batches)]
-
-        k,v = list(self.observed_variables.items())[0]
-        for k,v in self.observed_variables.items():
-            if k in self.data and v.is_datamodel:
-                batches_k = np.split(data[k], num_batches, axis=0)
-                for i in range(num_batches):
-                    self.batches[i][k] = batches_k[i]
-            elif k in self.data:
-                for i in range(num_batches):
-                    self.batches[i][k] = data[k]
-
-
-
     @flatten_result
     @util.tf_run_ignored
     def log_prob(self):
         """ Computes the log probabilities of a (set of) sample(s)"""
+        with util.interceptor.enable_interceptor(*self.enable_interceptor_variables):
+            with contextmanager.observe(self.observed_variables, self.data):
+                result = util.runtime.try_run({k: v.log_prob(v.value) for k, v in self.target_variables.items()})
 
-        results = []
-        for batch in self.batches:
-            with util.interceptor.enable_interceptor(*self.enable_interceptor_variables):
-                with contextmanager.observe(self.observed_variables, batch):
-                    result = util.runtime.try_run({k: v.log_prob(v.value) for k, v in self.target_variables.items()})
-
-            results.append(result)
-
-
-        return {k:np.vstack([r[k] for r in results])[:self.data_size] for k in self.target_variables.keys()}
+        return result
 
     def sum_log_prob(self):
         """ Computes the sum of the log probabilities (evaluated) of a (set of) sample(s)"""
@@ -92,26 +54,19 @@ def sum_log_prob(self):
     @util.tf_run_ignored
     def sample(self, size=1):
         """ Generates a sample for eache variable in the model """
+        with util.interceptor.enable_interceptor(*self.enable_interceptor_variables):
+            with contextmanager.observe(self.observed_variables, self.data):
+                # each iteration for `size` run the dict in the session, so if there are dependencies among random vars
+                # they are computed in the same graph operations, and reflected in the results
+                samples = [util.runtime.try_run(self.target_variables) for _ in range(size)]
+
+        if size == 1:
+            result = samples[0]
+        else:
+            # compact all samples in one single dict
+            result = {k: np.array([sample[k] for sample in samples]) for k in self.target_variables.keys()}
 
-        results = []
-        for batch in self.batches:
-            with util.interceptor.enable_interceptor(*self.enable_interceptor_variables):
-                with contextmanager.observe(self.observed_variables, batch):
-                    # each iteration for `size` run the dict in the session, so if there are dependencies among random vars
-                    # they are computed in the same graph operations, and reflected in the results
-                    samples = [util.runtime.try_run(self.target_variables) for _ in range(size)]
-
-            if size == 1:
-                result = samples[0]
-            else:
-                # compact all samples in one single dict
-                result = {k: np.array([sample[k] for sample in samples]) for k in self.target_variables.keys()}
-
-            results.append(result)
-
-
-        return {k:np.vstack([r[k] for r in results])[:self.data_size] for k in self.target_variables.keys()}
-
+        return result
 
     @flatten_result
     @util.tf_run_ignored
@@ -149,27 +104,8 @@ def filter_parameters(varname, parameters):
 
             return {k: util.runtime.try_run(v) for k, v in parameters.items() if k in selected_parameters}
 
-        # function that merges the parameters of 2 batches
-        def merge_params(p1, p2):
-            out = {}
-            for k,v in p1.items():
-                if np.ndim(v) == 0:
-                    out[k] = v
-                else:
-                    out[k] = np.vstack([p1[k], p2[k]])[0:self.data_size]
-            return out
-
-        # get the parameter for each batch
-        result = {}
-
-        for batch in self.batches:
-            with contextmanager.observe(self.observed_variables, batch):
-                r = {k: filter_parameters(k, v.parameters)
-                          for k, v in self.target_variables.items()}
-
-            if len(result.keys()) == 0:
-                result = r
-            else:
-                result = {var:merge_params(result[var], r[var]) for var in self.target_variables.keys()}
+        with contextmanager.observe(self.observed_variables, self.data):
+            result = {k: filter_parameters(k, v.parameters)
+                      for k, v in self.target_variables.items()}
 
         return result
diff --git a/inferpy/util/common.py b/inferpy/util/common.py
index c6e9e98..63d2755 100644
--- a/inferpy/util/common.py
+++ b/inferpy/util/common.py
@@ -1,39 +1,37 @@
-""" Obtained from Keras GitHub repository: https://github.com/keras-team/keras/blob/master/keras/backend/common.py
-
+"""
+Obtained from Keras GitHub repository: https://github.com/keras-team/keras/blob/master/keras/backend/common.py
 """
 
 _FLOATX = "float32"
 
 
 def floatx():
-    """Returns the default float type, as a string.  (e.g. float16, float32, float64).
-
-    Returns:
-        String: the current default float type.
-
-    Example:
+    """Returns the default float type, as a string.
+    (e.g. 'float16', 'float32', 'float64').
+    # Returns
+        String, the current default float type.
+    # Example
+    ```python
         >>> inf.floatx()
         'float32'
-
-
+    ```
     """
     return _FLOATX
 
 
 def set_floatx(floatx):
-    """ Sets the default float type.
-
-    Args:
+    """Sets the default float type.
+    # Arguments
         floatx: String, 'float16', 'float32', or 'float64'.
-
-    Example:
+    # Example
+    ```python
         >>> from keras import backend as K
         >>> inf.floatx()
         'float32'
         >>> inf.set_floatx('float16')
         >>> inf..floatx()
         'float16'
-
+    ```
     """
     global _FLOATX
     if floatx not in {'float16', 'float32', 'float64'}:
diff --git a/inferpy/util/iterables.py b/inferpy/util/iterables.py
index b59df78..c1d1cbe 100644
--- a/inferpy/util/iterables.py
+++ b/inferpy/util/iterables.py
@@ -1,18 +1,13 @@
 
 def get_shape(x):
-    """ Get the shape of an element x. If it is an element with a shape attribute, return it. If it is a list with more than
+    """
+    Get the shape of an element x. If it is an element with a shape attribute, return it. If it is a list with more than
     one element, compute the shape by checking the len, and the shape of internal elements. In that case, the shape must
     be consistent. Finally, in other case return () as shape.
 
-    Args:
-        x: The element to compute its shape
-
-    Raises:
-        class `ValueError`: list shape not consistent
-
-    Returns:
-        A tuple with the shape of `x`
-
+    :param x: The element to compute its shape
+    :raises : class `ValueError`: list shape not consistent
+    :returns: A tuple with the shape of `x`
     """
     if isinstance(x, list) and len(x) > 0:
         shapes = [get_shape(subx) for subx in x]
diff --git a/tests/test_queries.py b/tests/test_queries.py
new file mode 100644
index 0000000..3f8d29d
--- /dev/null
+++ b/tests/test_queries.py
@@ -0,0 +1,178 @@
+import numpy as np
+import pytest
+import tensorflow as tf
+
+import inferpy as inf
+from inferpy import util
+from tests import no_raised_exc
+
+
+def init_model():
+    @inf.probmodel
+    def linear_reg(d):
+        w0 = inf.Normal(0, 1, name="w0")
+        w = inf.Normal(np.zeros([d, 1]), 1, name="w")
+
+        with inf.datamodel():
+            x = inf.Normal(tf.ones(d), 2, name="x")
+            y = inf.Normal(w0 + x @ w, 1.0, name="y")
+
+
+    @inf.probmodel
+    def qmodel(d):
+        qw0_loc = inf.Parameter(0., name="qw0_loc")
+        qw0_scale = tf.math.softplus(inf.Parameter(1., name="qw0_scale"))
+        qw0 = inf.Normal(qw0_loc, qw0_scale, name="w0")
+
+        qw_loc = inf.Parameter(np.zeros([d, 1]), name="qw_loc")
+        qw_scale = tf.math.softplus(inf.Parameter(tf.ones([d, 1]), name="qw_scale"))
+        qw = inf.Normal(qw_loc, qw_scale, name="w")
+
+
+    # create an instance of the model
+    m = linear_reg(d=2)
+    q = qmodel(2)
+    # create toy data
+    N = 1000
+    data = m.prior(["x", "y"], data={"w0": 0, "w": [[2], [1]]}, size_datamodel=N).sample()
+
+
+    x_train = data["x"]
+    y_train = data["y"]
+
+    # set and run the inference
+    VI = inf.inference.VI(qmodel(2), epochs=500)
+    m.fit({"x": x_train, "y": y_train}, VI)
+
+    return m, x_train, y_train
+
+
+
+def test_prior():
+    m, x_train, y_train = init_model()
+
+    s = m.prior().sample()
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["w0"].shape == ()
+    assert s["w"].shape == (2, 1)
+    assert s["x"].shape == (1, 2)
+    assert s["y"].shape == (1, 1)
+
+
+def test_prior_size():
+    m, x_train, y_train = init_model()
+    s = m.prior().sample(size=5)
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["w0"].shape == (5,)
+    assert s["w"].shape == (5, 2, 1)
+    assert s["x"].shape == (5, 1, 2)
+    assert s["y"].shape == (5, 1, 1)
+
+
+def test_prior_size_datamodel():
+    m, x_train, y_train = init_model()
+    s = m.prior(size_datamodel=10).sample()
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["w0"].shape == ()
+    assert s["w"].shape == (2, 1)
+    assert s["x"].shape == (10, 2)
+    assert s["y"].shape == (10, 1)
+
+def test_prior_size2():
+    m, x_train, y_train = init_model()
+    s = m.prior(size_datamodel=10).sample(size=5)
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["w0"].shape == (5,)
+    assert s["w"].shape == (5, 2, 1)
+    assert s["x"].shape == (5, 10, 2)
+    assert s["y"].shape == (5, 10, 1)
+
+
+def test_prior_data():
+    m, x_train, y_train = init_model()
+    s = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((10, 1))}).sample(size=5)
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["w0"].shape == (5,)
+    assert s["w"].shape == (5, 2, 1)
+    assert s["x"].shape == (5, 10, 2)
+    assert s["y"].shape == (5, 10, 1)
+
+
+def test_prior_batches():
+    m, x_train, y_train = init_model()
+    s = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((15, 1))}).sample(size=5)  # ERROR
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["w0"].shape == (5,)
+    assert s["w"].shape == (5, 2, 1)
+    assert s["x"].shape == (5, 15, 2)
+    assert s["y"].shape == (5, 15, 1)
+
+
+def test_post():
+    m, x_train, y_train = init_model()
+    s = m.posterior_predictive().sample()
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["x"].shape == (1000, 2)
+    assert s["y"].shape == (1000, 1)
+
+
+def test_post_size():
+    m, x_train, y_train = init_model()
+    s = m.posterior_predictive().sample(size=5)
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["x"].shape == (5, 1000, 2)
+    assert s["y"].shape == (5, 1000, 1)
+
+
+def test_post_data():
+    m, x_train, y_train = init_model()
+    s = m.posterior_predictive(data={"x":x_train}).sample(3)
+
+    print({k:v.shape for k,v in s.items()})
+
+    assert s["x"].shape == (3,1000,2)
+    assert s["y"].shape == (3,1000,1)
+
+
+
+def test_post_data2():
+    m, x_train, y_train = init_model()
+    s = m.posterior_predictive(data={"w0": 0, "w": [[2], [1]], "y": np.zeros((1000, 1))}).sample(3)
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["x"].shape == (3, 1000, 2)
+    assert s["y"].shape == (3, 1000, 1)
+
+
+def test_post_batches():
+    m, x_train, y_train = init_model()
+    s = m.posterior_predictive(data={"y": np.zeros((1300, 1))}).sample(3)  # ERROR
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["x"].shape == (3, 1300, 2)
+    assert s["y"].shape == (3, 1300, 1)
+
+
+
+
+
+
+

From 44fe6dc0db96462171b98582017d6d6e9b57819f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Caba=C3=B1as?= <rcabanas@idsia.ch>
Date: Thu, 15 Oct 2020 10:16:20 +0200
Subject: [PATCH 04/11] test queries

---
 tests/test_queries.py | 178 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 178 insertions(+)
 create mode 100644 tests/test_queries.py

diff --git a/tests/test_queries.py b/tests/test_queries.py
new file mode 100644
index 0000000..3f8d29d
--- /dev/null
+++ b/tests/test_queries.py
@@ -0,0 +1,178 @@
+import numpy as np
+import pytest
+import tensorflow as tf
+
+import inferpy as inf
+from inferpy import util
+from tests import no_raised_exc
+
+
+def init_model():
+    @inf.probmodel
+    def linear_reg(d):
+        w0 = inf.Normal(0, 1, name="w0")
+        w = inf.Normal(np.zeros([d, 1]), 1, name="w")
+
+        with inf.datamodel():
+            x = inf.Normal(tf.ones(d), 2, name="x")
+            y = inf.Normal(w0 + x @ w, 1.0, name="y")
+
+
+    @inf.probmodel
+    def qmodel(d):
+        qw0_loc = inf.Parameter(0., name="qw0_loc")
+        qw0_scale = tf.math.softplus(inf.Parameter(1., name="qw0_scale"))
+        qw0 = inf.Normal(qw0_loc, qw0_scale, name="w0")
+
+        qw_loc = inf.Parameter(np.zeros([d, 1]), name="qw_loc")
+        qw_scale = tf.math.softplus(inf.Parameter(tf.ones([d, 1]), name="qw_scale"))
+        qw = inf.Normal(qw_loc, qw_scale, name="w")
+
+
+    # create an instance of the model
+    m = linear_reg(d=2)
+    q = qmodel(2)
+    # create toy data
+    N = 1000
+    data = m.prior(["x", "y"], data={"w0": 0, "w": [[2], [1]]}, size_datamodel=N).sample()
+
+
+    x_train = data["x"]
+    y_train = data["y"]
+
+    # set and run the inference
+    VI = inf.inference.VI(qmodel(2), epochs=500)
+    m.fit({"x": x_train, "y": y_train}, VI)
+
+    return m, x_train, y_train
+
+
+
+def test_prior():
+    m, x_train, y_train = init_model()
+
+    s = m.prior().sample()
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["w0"].shape == ()
+    assert s["w"].shape == (2, 1)
+    assert s["x"].shape == (1, 2)
+    assert s["y"].shape == (1, 1)
+
+
+def test_prior_size():
+    m, x_train, y_train = init_model()
+    s = m.prior().sample(size=5)
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["w0"].shape == (5,)
+    assert s["w"].shape == (5, 2, 1)
+    assert s["x"].shape == (5, 1, 2)
+    assert s["y"].shape == (5, 1, 1)
+
+
+def test_prior_size_datamodel():
+    m, x_train, y_train = init_model()
+    s = m.prior(size_datamodel=10).sample()
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["w0"].shape == ()
+    assert s["w"].shape == (2, 1)
+    assert s["x"].shape == (10, 2)
+    assert s["y"].shape == (10, 1)
+
+def test_prior_size2():
+    m, x_train, y_train = init_model()
+    s = m.prior(size_datamodel=10).sample(size=5)
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["w0"].shape == (5,)
+    assert s["w"].shape == (5, 2, 1)
+    assert s["x"].shape == (5, 10, 2)
+    assert s["y"].shape == (5, 10, 1)
+
+
+def test_prior_data():
+    m, x_train, y_train = init_model()
+    s = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((10, 1))}).sample(size=5)
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["w0"].shape == (5,)
+    assert s["w"].shape == (5, 2, 1)
+    assert s["x"].shape == (5, 10, 2)
+    assert s["y"].shape == (5, 10, 1)
+
+
+def test_prior_batches():
+    m, x_train, y_train = init_model()
+    s = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((15, 1))}).sample(size=5)  # ERROR
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["w0"].shape == (5,)
+    assert s["w"].shape == (5, 2, 1)
+    assert s["x"].shape == (5, 15, 2)
+    assert s["y"].shape == (5, 15, 1)
+
+
+def test_post():
+    m, x_train, y_train = init_model()
+    s = m.posterior_predictive().sample()
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["x"].shape == (1000, 2)
+    assert s["y"].shape == (1000, 1)
+
+
+def test_post_size():
+    m, x_train, y_train = init_model()
+    s = m.posterior_predictive().sample(size=5)
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["x"].shape == (5, 1000, 2)
+    assert s["y"].shape == (5, 1000, 1)
+
+
+def test_post_data():
+    m, x_train, y_train = init_model()
+    s = m.posterior_predictive(data={"x":x_train}).sample(3)
+
+    print({k:v.shape for k,v in s.items()})
+
+    assert s["x"].shape == (3,1000,2)
+    assert s["y"].shape == (3,1000,1)
+
+
+
+def test_post_data2():
+    m, x_train, y_train = init_model()
+    s = m.posterior_predictive(data={"w0": 0, "w": [[2], [1]], "y": np.zeros((1000, 1))}).sample(3)
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["x"].shape == (3, 1000, 2)
+    assert s["y"].shape == (3, 1000, 1)
+
+
+def test_post_batches():
+    m, x_train, y_train = init_model()
+    s = m.posterior_predictive(data={"y": np.zeros((1300, 1))}).sample(3)  # ERROR
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["x"].shape == (3, 1300, 2)
+    assert s["y"].shape == (3, 1300, 1)
+
+
+
+
+
+
+

From 0d4eaaa70d3220a55e5f4c8cb6777ad2ca243d4c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Caba=C3=B1as?= <rcabanas@idsia.ch>
Date: Thu, 15 Oct 2020 11:46:59 +0200
Subject: [PATCH 05/11] query sample fixed

---
 inferpy/queries/query.py | 15 ++++++++++++---
 tests/test_queries.py    |  8 +++++++-
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/inferpy/queries/query.py b/inferpy/queries/query.py
index 8d64c3f..379a6f9 100644
--- a/inferpy/queries/query.py
+++ b/inferpy/queries/query.py
@@ -37,9 +37,15 @@ def __init__(self, variables, target_names=None, data={}, enable_interceptor_var
         self.data = data
         self.enable_interceptor_variables = enable_interceptor_variables
 
-
+        # plateau size
         self.batch_size = [v for (k,v) in self.observed_variables.items() if v.is_datamodel][0].shape[0].value
-        self.data_size = self.data[[k for (k,v) in self.observed_variables.items() if k in self.data.keys() and v.is_datamodel][0]].shape[0]
+
+        # size of the observed variables in data
+        obs_datamodel = [k for (k,v) in self.observed_variables.items() if k in self.data.keys() and v.is_datamodel]
+        if len(obs_datamodel)>0:
+            self.data_size = self.data[obs_datamodel[0]].shape[0]
+        else:
+            self.data_size = self.batch_size
 
         ####
 
@@ -80,7 +86,8 @@ def log_prob(self):
 
             results.append(result)
 
-
+        if len(self.batches)==1:
+            return results[0]
         return {k:np.vstack([r[k] for r in results])[:self.data_size] for k in self.target_variables.keys()}
 
     def sum_log_prob(self):
@@ -109,6 +116,8 @@ def sample(self, size=1):
 
             results.append(result)
 
+        if len(self.batches)==1:
+            return results[0]
 
         return {k:np.vstack([r[k] for r in results])[:self.data_size] for k in self.target_variables.keys()}
 
diff --git a/tests/test_queries.py b/tests/test_queries.py
index 3f8d29d..f647397 100644
--- a/tests/test_queries.py
+++ b/tests/test_queries.py
@@ -52,6 +52,7 @@ def test_prior():
     m, x_train, y_train = init_model()
 
     s = m.prior().sample()
+    l = m.prior().log_prob(s)
 
     print({k: v.shape for k, v in s.items()})
 
@@ -60,6 +61,11 @@ def test_prior():
     assert s["x"].shape == (1, 2)
     assert s["y"].shape == (1, 1)
 
+    assert l["w0"].shape == ()
+    assert l["w"].shape == (2, 1)
+    assert l["x"].shape == (1, 2)
+    assert l["y"].shape == (1, 1)
+
 
 def test_prior_size():
     m, x_train, y_train = init_model()
@@ -153,7 +159,7 @@ def test_post_data():
 
 def test_post_data2():
     m, x_train, y_train = init_model()
-    s = m.posterior_predictive(data={"w0": 0, "w": [[2], [1]], "y": np.zeros((1000, 1))}).sample(3)
+    s = m.posterior_predictive(data={"w0": 0, "w": [[2], [1]], "x": np.zeros((1000, 2))}).sample(3)
 
     print({k: v.shape for k, v in s.items()})
 

From aa9cbdb0b91de83e5bc1869b2094d36884649712 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Caba=C3=B1as?= <rcabanas@idsia.ch>
Date: Thu, 15 Oct 2020 11:47:27 +0200
Subject: [PATCH 06/11] query sample fixed

---
 tests/test_queries.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tests/test_queries.py b/tests/test_queries.py
index f647397..db664fa 100644
--- a/tests/test_queries.py
+++ b/tests/test_queries.py
@@ -52,7 +52,7 @@ def test_prior():
     m, x_train, y_train = init_model()
 
     s = m.prior().sample()
-    l = m.prior().log_prob(s)
+    # l = m.prior().log_prob(s)
 
     print({k: v.shape for k, v in s.items()})
 
@@ -61,10 +61,10 @@ def test_prior():
     assert s["x"].shape == (1, 2)
     assert s["y"].shape == (1, 1)
 
-    assert l["w0"].shape == ()
-    assert l["w"].shape == (2, 1)
-    assert l["x"].shape == (1, 2)
-    assert l["y"].shape == (1, 1)
+    # assert l["w0"].shape == ()
+    # assert l["w"].shape == (2, 1)
+    # assert l["x"].shape == (1, 2)
+    # assert l["y"].shape == (1, 1)
 
 
 def test_prior_size():

From 3f43934763177b0f1c025be62c709e3d13c9fefd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Caba=C3=B1as?= <rcabanas@idsia.ch>
Date: Thu, 15 Oct 2020 15:59:38 +0200
Subject: [PATCH 07/11] query sample fixed2

---
 inferpy/queries/query.py | 27 ++++++++----
 tests/test_queries.py    | 88 ++++++++++++++++++++++++++++++++++++----
 2 files changed, 100 insertions(+), 15 deletions(-)

diff --git a/inferpy/queries/query.py b/inferpy/queries/query.py
index 379a6f9..2bf3398 100644
--- a/inferpy/queries/query.py
+++ b/inferpy/queries/query.py
@@ -86,20 +86,19 @@ def log_prob(self):
 
             results.append(result)
 
-        if len(self.batches)==1:
-            return results[0]
-        return {k:np.vstack([r[k] for r in results])[:self.data_size] for k in self.target_variables.keys()}
+        return self.process_output(results)
+
 
     def sum_log_prob(self):
         """ Computes the sum of the log probabilities (evaluated) of a (set of) sample(s)"""
         # The decorator is not needed here because this function returns a single value
         return np.sum([np.mean(lp) for lp in self.log_prob(simplify_result=False).values()])
 
+
     @flatten_result
     @util.tf_run_ignored
     def sample(self, size=1):
         """ Generates a sample for eache variable in the model """
-
         results = []
         for batch in self.batches:
             with util.interceptor.enable_interceptor(*self.enable_interceptor_variables):
@@ -116,10 +115,24 @@ def sample(self, size=1):
 
             results.append(result)
 
-        if len(self.batches)==1:
-            return results[0]
+        return self.process_output(results)
+
+
+    def process_output(self, results):
+        out = {}
+
+        for r in results:
+            for k,v in self.target_variables.items():
+                if k not in out:
+                    out[k] = r[k]
+                elif v.is_datamodel:
+                    # merge batches
+                    out[k] = np.concatenate([out[k], r[k]], axis=-2)
+                    if self.data_size<out[k].shape[-2]:
+                        out[k] = np.take(out[k], range(self.data_size), axis=-2)
+
+        return out
 
-        return {k:np.vstack([r[k] for r in results])[:self.data_size] for k in self.target_variables.keys()}
 
 
     @flatten_result
diff --git a/tests/test_queries.py b/tests/test_queries.py
index db664fa..8dcae3c 100644
--- a/tests/test_queries.py
+++ b/tests/test_queries.py
@@ -52,7 +52,7 @@ def test_prior():
     m, x_train, y_train = init_model()
 
     s = m.prior().sample()
-    # l = m.prior().log_prob(s)
+    l = m.prior().log_prob()
 
     print({k: v.shape for k, v in s.items()})
 
@@ -61,15 +61,17 @@ def test_prior():
     assert s["x"].shape == (1, 2)
     assert s["y"].shape == (1, 1)
 
-    # assert l["w0"].shape == ()
-    # assert l["w"].shape == (2, 1)
-    # assert l["x"].shape == (1, 2)
-    # assert l["y"].shape == (1, 1)
+    assert l["w0"].shape == ()
+    assert l["w"].shape == (2, 1)
+    assert l["x"].shape == (1, 2)
+    assert l["y"].shape == (1, 1)
 
 
 def test_prior_size():
     m, x_train, y_train = init_model()
     s = m.prior().sample(size=5)
+    l = m.prior().sample(size=5)
+
 
     print({k: v.shape for k, v in s.items()})
 
@@ -78,10 +80,17 @@ def test_prior_size():
     assert s["x"].shape == (5, 1, 2)
     assert s["y"].shape == (5, 1, 1)
 
+    assert l["w0"].shape == (5,)
+    assert l["w"].shape == (5, 2, 1)
+    assert l["x"].shape == (5, 1, 2)
+    assert l["y"].shape == (5, 1, 1)
+
 
 def test_prior_size_datamodel():
     m, x_train, y_train = init_model()
     s = m.prior(size_datamodel=10).sample()
+    l = m.prior(size_datamodel=10).log_prob()
+
 
     print({k: v.shape for k, v in s.items()})
 
@@ -90,6 +99,12 @@ def test_prior_size_datamodel():
     assert s["x"].shape == (10, 2)
     assert s["y"].shape == (10, 1)
 
+    assert l["w0"].shape == ()
+    assert l["w"].shape == (2, 1)
+    assert l["x"].shape == (10, 2)
+    assert l["y"].shape == (10, 1)
+
+
 def test_prior_size2():
     m, x_train, y_train = init_model()
     s = m.prior(size_datamodel=10).sample(size=5)
@@ -101,22 +116,50 @@ def test_prior_size2():
     assert s["x"].shape == (5, 10, 2)
     assert s["y"].shape == (5, 10, 1)
 
-
 def test_prior_data():
     m, x_train, y_train = init_model()
     s = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((10, 1))}).sample(size=5)
+    l = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((10, 1))}).log_prob()
 
     print({k: v.shape for k, v in s.items()})
+    print({k: v.shape for k, v in l.items()})
 
     assert s["w0"].shape == (5,)
     assert s["w"].shape == (5, 2, 1)
     assert s["x"].shape == (5, 10, 2)
     assert s["y"].shape == (5, 10, 1)
 
+    assert l["w0"].shape == ()
+    assert l["w"].shape == (2, 1)
+    assert l["x"].shape == (10, 2)
+    assert l["y"].shape == (10, 1)
+
 
 def test_prior_batches():
     m, x_train, y_train = init_model()
-    s = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((15, 1))}).sample(size=5)  # ERROR
+    s = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((15, 1))}).sample()
+    l = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((15, 1))}).log_prob()
+
+
+    print({k: v.shape for k, v in s.items()})
+
+    assert s["w0"].shape == ()
+    assert s["w"].shape == (2, 1)
+    assert s["x"].shape == (15, 2)
+    assert s["y"].shape == (15, 1)
+
+    #
+    assert l["w0"].shape == ()
+    assert l["w"].shape == (2, 1)
+    assert l["x"].shape == (15, 2)
+    assert l["y"].shape == (15, 1)
+
+
+def test_prior_batches2():
+    m, x_train, y_train = init_model()
+    s = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((15, 1))}).sample(size=5)
+    l = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((15, 1))}).log_prob()
+
 
     print({k: v.shape for k, v in s.items()})
 
@@ -125,57 +168,86 @@ def test_prior_batches():
     assert s["x"].shape == (5, 15, 2)
     assert s["y"].shape == (5, 15, 1)
 
+    assert l["w0"].shape == ()
+    assert l["w"].shape == (2, 1)
+    assert l["x"].shape == (15, 2)
+    assert l["y"].shape == (15, 1)
 
 def test_post():
     m, x_train, y_train = init_model()
     s = m.posterior_predictive().sample()
+    l = m.posterior_predictive().log_prob()
+
 
     print({k: v.shape for k, v in s.items()})
 
     assert s["x"].shape == (1000, 2)
     assert s["y"].shape == (1000, 1)
 
+    assert l["x"].shape == (1000, 2)
+    assert l["y"].shape == (1000, 1)
+
 
 def test_post_size():
     m, x_train, y_train = init_model()
     s = m.posterior_predictive().sample(size=5)
+    l = m.posterior_predictive().log_prob()
+
 
     print({k: v.shape for k, v in s.items()})
 
     assert s["x"].shape == (5, 1000, 2)
     assert s["y"].shape == (5, 1000, 1)
 
+    assert l["x"].shape == (1000, 2)
+    assert l["y"].shape == (1000, 1)
+
 
 def test_post_data():
     m, x_train, y_train = init_model()
     s = m.posterior_predictive(data={"x":x_train}).sample(3)
+    l = m.posterior_predictive(data={"x":x_train}).log_prob()
+
 
     print({k:v.shape for k,v in s.items()})
 
     assert s["x"].shape == (3,1000,2)
     assert s["y"].shape == (3,1000,1)
 
+    assert l["x"].shape == (1000, 2)
+    assert l["y"].shape == (1000, 1)
+
 
 
 def test_post_data2():
     m, x_train, y_train = init_model()
     s = m.posterior_predictive(data={"w0": 0, "w": [[2], [1]], "x": np.zeros((1000, 2))}).sample(3)
+    l = m.posterior_predictive(data={"w0": 0, "w": [[2], [1]], "x": np.zeros((1000, 2))}).log_prob()
+
 
     print({k: v.shape for k, v in s.items()})
 
     assert s["x"].shape == (3, 1000, 2)
     assert s["y"].shape == (3, 1000, 1)
 
+    assert l["x"].shape == (1000, 2)
+    assert l["y"].shape == (1000, 1)
+
+
 
 def test_post_batches():
     m, x_train, y_train = init_model()
-    s = m.posterior_predictive(data={"y": np.zeros((1300, 1))}).sample(3)  # ERROR
+    s = m.posterior_predictive(data={"y": np.zeros((1300, 1))}).sample(3)
+    l = m.posterior_predictive(data={"y": np.zeros((1300, 1))}).log_prob()
 
     print({k: v.shape for k, v in s.items()})
 
     assert s["x"].shape == (3, 1300, 2)
     assert s["y"].shape == (3, 1300, 1)
 
+    assert l["x"].shape == (1300, 2)
+    assert l["y"].shape == (1300, 1)
+
 
 
 

From 59aeb5995c18a61a4eecc3784a2edba9abcbbcfb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Caba=C3=B1as?= <rcabanas@idsia.ch>
Date: Thu, 15 Oct 2020 17:05:41 +0200
Subject: [PATCH 08/11] query sample fixed3

---
 inferpy/inference/mcmc.py | 6 ++++++
 inferpy/queries/query.py  | 3 ++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/inferpy/inference/mcmc.py b/inferpy/inference/mcmc.py
index a5f1f0b..65a77bc 100644
--- a/inferpy/inference/mcmc.py
+++ b/inferpy/inference/mcmc.py
@@ -83,6 +83,12 @@ def update(self, data):
         self.states = {name: models.Empirical(states, event_ndims=len(states.shape) - 1, name=name)
                        for name, states in zip(self.hiddenvars_name, variables_states)}
 
+
+        # set the is_datamodel flag to the results
+        for k,v in self.expanded_variables.items():
+            if k in self.states:
+                self.states[k].is_datamodel = v.is_datamodel
+
     def posterior(self, target_names=None, data={}):
         return Query(self.states, target_names, data)
 
diff --git a/inferpy/queries/query.py b/inferpy/queries/query.py
index 2bf3398..7a1ab48 100644
--- a/inferpy/queries/query.py
+++ b/inferpy/queries/query.py
@@ -38,7 +38,8 @@ def __init__(self, variables, target_names=None, data={}, enable_interceptor_var
         self.enable_interceptor_variables = enable_interceptor_variables
 
         # plateau size
-        self.batch_size = [v for (k,v) in self.observed_variables.items() if v.is_datamodel][0].shape[0].value
+        vars_datamodel = [v for (k,v) in self.observed_variables.items() if v.is_datamodel]
+        self.batch_size = vars_datamodel[0].shape[0].value if len(vars_datamodel)>0 else 1
 
         # size of the observed variables in data
         obs_datamodel = [k for (k,v) in self.observed_variables.items() if k in self.data.keys() and v.is_datamodel]

From 99a30100913ef698c076e068c42011346181292d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Caba=C3=B1as?= <rcabanas@idsia.ch>
Date: Wed, 21 Oct 2020 17:01:27 +0200
Subject: [PATCH 09/11] fixed queries parameters

---
 inferpy/queries/query.py |  7 ++--
 tests/test_queries.py    | 71 ++++++++++++++++++++++++++++++++++++----
 2 files changed, 68 insertions(+), 10 deletions(-)

diff --git a/inferpy/queries/query.py b/inferpy/queries/query.py
index 7a1ab48..37375bb 100644
--- a/inferpy/queries/query.py
+++ b/inferpy/queries/query.py
@@ -121,7 +121,6 @@ def sample(self, size=1):
 
     def process_output(self, results):
         out = {}
-
         for r in results:
             for k,v in self.target_variables.items():
                 if k not in out:
@@ -175,8 +174,11 @@ def filter_parameters(varname, parameters):
         # function that merges the parameters of 2 batches
         def merge_params(p1, p2):
             out = {}
+            var = self.observed_variables[p1["name"]]
             for k,v in p1.items():
-                if np.ndim(v) == 0:
+                if np.ndim(v) == 0 or  \
+                        (var.is_datamodel and len(var.sample_shape.as_list())>0) or \
+                        not var.is_datamodel:
                     out[k] = v
                 else:
                     out[k] = np.vstack([p1[k], p2[k]])[0:self.data_size]
@@ -194,5 +196,4 @@ def merge_params(p1, p2):
                 result = r
             else:
                 result = {var:merge_params(result[var], r[var]) for var in self.target_variables.keys()}
-
         return result
diff --git a/tests/test_queries.py b/tests/test_queries.py
index 8dcae3c..8df6532 100644
--- a/tests/test_queries.py
+++ b/tests/test_queries.py
@@ -53,8 +53,11 @@ def test_prior():
 
     s = m.prior().sample()
     l = m.prior().log_prob()
+    p = m.prior().parameters()
 
     print({k: v.shape for k, v in s.items()})
+    print({k: v.shape for k, v in l.items()})
+    print({k: np.shape(p[k]["loc"]) for k, v in p.items()})
 
     assert s["w0"].shape == ()
     assert s["w"].shape == (2, 1)
@@ -66,14 +69,19 @@ def test_prior():
     assert l["x"].shape == (1, 2)
     assert l["y"].shape == (1, 1)
 
+    assert np.shape(p["w0"]["loc"]) == ()
+    assert np.shape(p["w"]["loc"]) == (2, 1)
+    assert np.shape(p["x"]["loc"]) == (2,)
+    assert np.shape(p["y"]["loc"]) == (1, 1)
+
 
 def test_prior_size():
     m, x_train, y_train = init_model()
     s = m.prior().sample(size=5)
     l = m.prior().sample(size=5)
 
-
     print({k: v.shape for k, v in s.items()})
+    print({k: v.shape for k, v in l.items()})
 
     assert s["w0"].shape == (5,)
     assert s["w"].shape == (5, 2, 1)
@@ -88,11 +96,16 @@ def test_prior_size():
 
 def test_prior_size_datamodel():
     m, x_train, y_train = init_model()
+
     s = m.prior(size_datamodel=10).sample()
     l = m.prior(size_datamodel=10).log_prob()
+    p = m.prior(size_datamodel=10).parameters()
 
 
     print({k: v.shape for k, v in s.items()})
+    print({k: v.shape for k, v in l.items()})
+    print({k: np.shape(p[k]["loc"]) for k, v in p.items()})
+
 
     assert s["w0"].shape == ()
     assert s["w"].shape == (2, 1)
@@ -104,6 +117,11 @@ def test_prior_size_datamodel():
     assert l["x"].shape == (10, 2)
     assert l["y"].shape == (10, 1)
 
+    assert np.shape(p["w0"]["loc"]) == ()
+    assert np.shape(p["w"]["loc"]) == (2, 1)
+    assert np.shape(p["x"]["loc"]) == (2,)
+    assert np.shape(p["y"]["loc"]) == (10, 1)
+
 
 def test_prior_size2():
     m, x_train, y_train = init_model()
@@ -120,9 +138,12 @@ def test_prior_data():
     m, x_train, y_train = init_model()
     s = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((10, 1))}).sample(size=5)
     l = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((10, 1))}).log_prob()
+    p = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((10, 1))}).parameters()
 
     print({k: v.shape for k, v in s.items()})
     print({k: v.shape for k, v in l.items()})
+    print({k: np.shape(p[k]["loc"]) for k, v in p.items()})
+
 
     assert s["w0"].shape == (5,)
     assert s["w"].shape == (5, 2, 1)
@@ -134,14 +155,21 @@ def test_prior_data():
     assert l["x"].shape == (10, 2)
     assert l["y"].shape == (10, 1)
 
+    assert np.shape(p["w0"]["loc"]) == ()
+    assert np.shape(p["w"]["loc"]) == (2, 1)
+    assert np.shape(p["x"]["loc"]) == (2,)
+    assert np.shape(p["y"]["loc"]) == (10, 1)
+
 
 def test_prior_batches():
     m, x_train, y_train = init_model()
     s = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((15, 1))}).sample()
     l = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((15, 1))}).log_prob()
-
+    p = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((15, 1))}).parameters()
 
     print({k: v.shape for k, v in s.items()})
+    print({k: v.shape for k, v in l.items()})
+    print({k: np.shape(p[k]["loc"]) for k, v in p.items()})
 
     assert s["w0"].shape == ()
     assert s["w"].shape == (2, 1)
@@ -154,14 +182,21 @@ def test_prior_batches():
     assert l["x"].shape == (15, 2)
     assert l["y"].shape == (15, 1)
 
+    assert np.shape(p["w0"]["loc"]) == ()
+    assert np.shape(p["w"]["loc"]) == (2, 1)
+    assert np.shape(p["x"]["loc"]) == (2,)
+    assert np.shape(p["y"]["loc"]) == (15, 1)
+
 
 def test_prior_batches2():
     m, x_train, y_train = init_model()
     s = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((15, 1))}).sample(size=5)
     l = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((15, 1))}).log_prob()
-
+    p = m.prior(size_datamodel=10, data={"w0": 0, "w": [[2], [1]], "y": np.zeros((15, 1))}).parameters()
 
     print({k: v.shape for k, v in s.items()})
+    print({k: v.shape for k, v in l.items()})
+    print({k: np.shape(p[k]["loc"]) for k, v in p.items()})
 
     assert s["w0"].shape == (5,)
     assert s["w"].shape == (5, 2, 1)
@@ -173,13 +208,21 @@ def test_prior_batches2():
     assert l["x"].shape == (15, 2)
     assert l["y"].shape == (15, 1)
 
+    assert np.shape(p["w0"]["loc"]) == ()
+    assert np.shape(p["w"]["loc"]) == (2, 1)
+    assert np.shape(p["x"]["loc"]) == (2,)
+    assert np.shape(p["y"]["loc"]) == (15, 1)
+
+
 def test_post():
     m, x_train, y_train = init_model()
     s = m.posterior_predictive().sample()
     l = m.posterior_predictive().log_prob()
-
+    p = m.posterior_predictive().parameters()
 
     print({k: v.shape for k, v in s.items()})
+    print({k: v.shape for k, v in l.items()})
+    print({k: np.shape(p[k]["loc"]) for k, v in p.items()})
 
     assert s["x"].shape == (1000, 2)
     assert s["y"].shape == (1000, 1)
@@ -187,6 +230,10 @@ def test_post():
     assert l["x"].shape == (1000, 2)
     assert l["y"].shape == (1000, 1)
 
+    assert np.shape(p["x"]["loc"]) == (2,)
+    assert np.shape(p["y"]["loc"]) == (1000, 1)
+
+
 
 def test_post_size():
     m, x_train, y_train = init_model()
@@ -207,9 +254,12 @@ def test_post_data():
     m, x_train, y_train = init_model()
     s = m.posterior_predictive(data={"x":x_train}).sample(3)
     l = m.posterior_predictive(data={"x":x_train}).log_prob()
+    p = m.posterior_predictive(data={"x":x_train}).parameters()
 
 
-    print({k:v.shape for k,v in s.items()})
+    print({k: v.shape for k, v in s.items()})
+    print({k: v.shape for k, v in l.items()})
+    print({k: np.shape(p[k]["loc"]) for k, v in p.items()})
 
     assert s["x"].shape == (3,1000,2)
     assert s["y"].shape == (3,1000,1)
@@ -217,6 +267,8 @@ def test_post_data():
     assert l["x"].shape == (1000, 2)
     assert l["y"].shape == (1000, 1)
 
+    assert np.shape(p["x"]["loc"]) == (2,)
+    assert np.shape(p["y"]["loc"]) == (1000, 1)
 
 
 def test_post_data2():
@@ -224,8 +276,9 @@ def test_post_data2():
     s = m.posterior_predictive(data={"w0": 0, "w": [[2], [1]], "x": np.zeros((1000, 2))}).sample(3)
     l = m.posterior_predictive(data={"w0": 0, "w": [[2], [1]], "x": np.zeros((1000, 2))}).log_prob()
 
-
     print({k: v.shape for k, v in s.items()})
+    print({k: v.shape for k, v in l.items()})
+    print({k: np.shape(p[k]["loc"]) for k, v in p.items()})
 
     assert s["x"].shape == (3, 1000, 2)
     assert s["y"].shape == (3, 1000, 1)
@@ -239,8 +292,11 @@ def test_post_batches():
     m, x_train, y_train = init_model()
     s = m.posterior_predictive(data={"y": np.zeros((1300, 1))}).sample(3)
     l = m.posterior_predictive(data={"y": np.zeros((1300, 1))}).log_prob()
+    p = m.posterior_predictive(data={"y": np.zeros((1300, 1))}).parameters()
 
     print({k: v.shape for k, v in s.items()})
+    print({k: v.shape for k, v in l.items()})
+    print({k: np.shape(p[k]["loc"]) for k, v in p.items()})
 
     assert s["x"].shape == (3, 1300, 2)
     assert s["y"].shape == (3, 1300, 1)
@@ -248,7 +304,8 @@ def test_post_batches():
     assert l["x"].shape == (1300, 2)
     assert l["y"].shape == (1300, 1)
 
-
+    assert np.shape(p["x"]["loc"]) == (2,)
+    assert np.shape(p["y"]["loc"]) == (1300, 1)
 
 
 

From c4c5e4bbb46a39b003876946cc3b7d1f0389cf21 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Caba=C3=B1as?= <rcabanas@idsia.ch>
Date: Wed, 21 Oct 2020 17:39:22 +0200
Subject: [PATCH 10/11] fixed queries parameters2

---
 tests/test_probmodel.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_probmodel.py b/tests/test_probmodel.py
index 25c41d0..959ba17 100644
--- a/tests/test_probmodel.py
+++ b/tests/test_probmodel.py
@@ -118,7 +118,7 @@ def model():
             inf.Normal(x, 1., name='y')
 
     N = 10
-    data_y = 1.0
+    data_y = np.array([1.0])
     m = model()
 
     sample_dict = m.prior(data={'y': data_y}).sample(N)

From 2d02767d89c298a324ea5a200d81d72fd19bd233 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Caba=C3=B1as?= <rcabanas@idsia.ch>
Date: Wed, 21 Oct 2020 17:42:22 +0200
Subject: [PATCH 11/11] fixed queries parameters3

---
 tests/test_queries.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/test_queries.py b/tests/test_queries.py
index 8df6532..4774513 100644
--- a/tests/test_queries.py
+++ b/tests/test_queries.py
@@ -278,7 +278,6 @@ def test_post_data2():
 
     print({k: v.shape for k, v in s.items()})
     print({k: v.shape for k, v in l.items()})
-    print({k: np.shape(p[k]["loc"]) for k, v in p.items()})
 
     assert s["x"].shape == (3, 1000, 2)
     assert s["y"].shape == (3, 1000, 1)