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/examples/probzoo/nl_digit_classifier.py b/examples/probzoo/nl_digit_classifier.py
new file mode 100644
index 0000000..97478a9
--- /dev/null
+++ b/examples/probzoo/nl_digit_classifier.py
@@ -0,0 +1,104 @@
+# 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
+
+
+@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=10000, 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
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/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 8697e97..37375bb 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,38 +37,103 @@ def __init__(self, variables, target_names=None, data={}, enable_interceptor_var
         self.data = data
         self.enable_interceptor_variables = enable_interceptor_variables
 
+        # plateau size
+        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]
+        if len(obs_datamodel)>0:
+            self.data_size = self.data[obs_datamodel[0]].shape[0]
+        else:
+            self.data_size = self.batch_size
+
+        ####
+
+
+        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 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 """
-        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 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 result
 
     @flatten_result
     @util.tf_run_ignored
@@ -104,8 +171,29 @@ 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 = {}
+            var = self.observed_variables[p1["name"]]
+            for k,v in p1.items():
+                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]
+            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
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_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)
diff --git a/tests/test_queries.py b/tests/test_queries.py
new file mode 100644
index 0000000..4774513
--- /dev/null
+++ b/tests/test_queries.py
@@ -0,0 +1,312 @@
+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()
+    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)
+    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 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)
+    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()
+    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)
+    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)
+
+    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()
+    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)
+    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)
+    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)
+
+    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)
+    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)
+
+    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)
+    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)
+
+    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)
+
+    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()
+    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()
+    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 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)
+
+    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():
+    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()})
+    print({k: v.shape for k, v in l.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)
+    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)
+
+    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)
+
+
+
+
+