diff --git a/docs/conf.py b/docs/conf.py
index edb0bbd..4534492 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -81,8 +81,8 @@
 
 # General information about the project.
 project = u'InferPy'
-copyright = u'2017, Javier Cózar, Rafael Cabañas, Antonio Salmerón, Andrés R. Masegosa'
-author = u'Javier Cózar, Rafael Cabañas, Antonio Salmerón, Andrés R. Masegosa'
+copyright = u'2017, Javier Cózar, Rafael Cabañas, Antonio Salmerón, Andrés R. Masegosa, Juan Jesús Ojeda'
+author = u'Javier Cózar, Rafael Cabañas, Antonio Salmerón, Andrés R. Masegosa, Juan Jesús Ojeda'
 
 # The version info for the project you're documenting, acts as replacement for
 # |version| and |release|, also used in various other places throughout the
diff --git a/docs/notes/contact.rst b/docs/notes/contact.rst
index 6f33008..f52da8a 100644
--- a/docs/notes/contact.rst
+++ b/docs/notes/contact.rst
@@ -5,7 +5,11 @@ Contact and Support
 If you have any question about the toolbox or if you want to collaborate in the project, please do not hesitate to
 contact us. You can do it through the following email address: inferpy.api@gmail.com
 
-
-For more technical questions, please use `Github issues <https://github.com/PGM-Lab/InferPy/issues>`_.
+- Javier Cózar (jcozar87@gmail.com)
+- Rafael Cabañas (rcabanas@idsia.ch)
+- Andrés R. Masegosa (andresmasegosa@ual.es)
+- Antonio Salmerón (antonio.salmeron@ual.es)
+- Juan Jesús Ojeda (juanje.ojeda@ual.es)
 
 
+For more technical questions, please use `Github issues <https://github.com/PGM-Lab/InferPy/issues>`_.
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/examples/probzoo/vae.py b/examples/probzoo/vae.py
index 75981b0..223eda3 100644
--- a/examples/probzoo/vae.py
+++ b/examples/probzoo/vae.py
@@ -68,7 +68,7 @@ def qmodel(k, d0, dx, encoder):
 q = qmodel(k, d0, dx, encoder)
 
 # set the inference algorithm
-SVI = inf.inference.SVI(q, epochs=5000)
+SVI = inf.inference.SVI(q, epochs=50)
 
 # learn the parameters
 m.fit({"x": x_train}, SVI)
@@ -83,5 +83,5 @@ def qmodel(k, d0, dx, encoder):
 m.posterior_predictive("x").sample(5)
 
 # decode values from the hidden representation
-m.posterior_predictive("x", data={"z": [2]}).sample(5)
+m.posterior_predictive("x", data={"z": 2})
 
diff --git a/examples/probzoo/vae_mnist_csvdata.py b/examples/probzoo/vae_mnist_csvdata.py
index ba5f8ee..2ecebd3 100644
--- a/examples/probzoo/vae_mnist_csvdata.py
+++ b/examples/probzoo/vae_mnist_csvdata.py
@@ -64,9 +64,9 @@ def qmodel(k, d0, dx, encoder):
 
 
 # create data loader from 10 csv files
-path = [f"./mnist_xtrain{i}.csv" for i in range(10)]
+path = [f"./datasets/mnist_xtrain{i}.csv" for i in range(10)]
 
-data_loader = CsvLoader(path, variables={"x" : range(dx)})
+data_loader = CsvLoader(path, var_dict={"x" : range(dx)})
 m.fit(data_loader.to_dict(), SVI)
 
 
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/data/preprocess.py b/inferpy/data/preprocess.py
new file mode 100644
index 0000000..c1d0bfc
--- /dev/null
+++ b/inferpy/data/preprocess.py
@@ -0,0 +1,124 @@
+import tensorflow as tf
+import numpy as np
+import math
+
+from inferpy.util.session import get_session
+
+def to_tensor(data, varnames=None):
+    varnames = varnames or data.keys()
+    return dict([
+        (k,tf.convert_to_tensor(v) if k in varnames else v)
+        for k,v in data.items()
+    ])
+
+
+def to_numpy(data, varnames=None):
+    varnames = varnames or data.keys()
+
+    # fn to convert a single value
+    def convert_to_numpy(x):
+        if isinstance(x, np.ndarray):
+             return x
+        if np.isscalar(x) or isinstance(x, list):
+            return np.array(x)
+        return convert_to_numpy(get_session().run(x))
+
+    # convert all the values
+    return dict([
+        (k,convert_to_numpy(v) if k in varnames else v)
+        for k,v in data.items()
+    ])
+
+
+
+
+
+
+def add_sample_dim(data, vars_datamodel):
+
+    data_out = dict()
+    for k,v in vars_datamodel.items():
+        if k in data:
+            data_out[k] = data[k]
+            # data cannot be an scalar nor list
+            if np.isscalar(data_out[k]) or isinstance(data_out[k], list):
+                data_out[k] = np.array(data_out[k])
+
+            # select the transforming gunctions depending on the type
+            if isinstance(data_out[k], tf.Tensor):
+                shape = data_out[k].get_shape().as_list()
+                expand_dims = tf.expand_dims
+                reshape = tf.reshape
+            elif isinstance(data_out[k], np.ndarray):
+                shape = data_out[k].shape
+                expand_dims = np.expand_dims
+                reshape = np.reshape
+            else:
+                raise ValueError("Unknown data type")
+
+            var_shape = vars_datamodel[k].shape.as_list()
+
+            # apply the transformation
+            if len(var_shape) - len(shape) == 1:
+                data_out[k] = expand_dims(data_out[k], axis=0)
+            elif len(var_shape) - len(shape) == 2:
+                data_out[k] = reshape(data_out[k], (1, 1))
+
+    return data_out
+
+
+def create_batches(data, vars_datamodel, data_size, batch_size, padding_last=True):
+    data_out = data.copy()
+
+    rows_to_add = round((math.ceil(data_size / batch_size) - data_size / batch_size) * batch_size)
+    num_batches = math.ceil(data_size / batch_size)
+
+    pad = {np.ndarray: np.pad, tf.Tensor:tf.pad}
+    split = {np.ndarray: np.split, tf.Tensor:tf.split}
+
+    # add padding if required
+    if padding_last:
+        for k, v in vars_datamodel.items():
+            if k in data_out:
+                padding = np.zeros([len(data_out[k].shape), 2], np.int32)
+                padding[0, 1] = rows_to_add
+                data_out[k] = pad[type(data_out[k])](data_out[k], padding, mode="constant")
+
+    # create batches
+    batches = [{} for _ in range(num_batches)]
+
+    for k in data_out.keys():
+        if k in vars_datamodel:
+            batches_k = split[type(data_out[k])](data_out[k], num_batches, axis=0)
+            for i in range(num_batches):
+                batches[i][k] = batches_k[i]
+        else:
+            for i in range(num_batches):
+                batches[i][k] = data_out[k]
+
+    return batches
+#
+#
+# data_ = data.copy()
+#
+#
+#
+#
+#
+# data_k.shape
+#
+#
+# ###
+# data = data_
+#
+# data["x"] = data["x"][0]
+#
+# vars_datamodel["z"].shape.as_list()
+#
+# data["x"] = 4
+#
+# data["x"].shape
+# add_sample_dim(to_tensor(data), vars_datamodel)
+#
+# data = to_tensor(data_)
+# data = data_
\ No newline at end of file
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..e96a2de 100644
--- a/inferpy/queries/query.py
+++ b/inferpy/queries/query.py
@@ -1,8 +1,14 @@
 import numpy as np
 import functools
 
+import math
+import tensorflow as tf
+
 from inferpy import contextmanager
 from inferpy import util
+from inferpy.data.preprocess import to_numpy, add_sample_dim, create_batches
+
+
 
 
 def flatten_result(f):
@@ -31,42 +37,99 @@ def __init__(self, variables, target_names=None, data={}, enable_interceptor_var
         self.target_variables = variables if not target_names else \
             {k: v for k, v in variables.items() if k in target_names}
 
-        self.observed_variables = variables
-        self.data = data
+        self.query_vars = variables
         self.enable_interceptor_variables = enable_interceptor_variables
 
+        vars_datamodel = dict([(k,v) for (k,v) in self.query_vars.items() if v.is_datamodel])
+
+        # all observations from vars in the data model must have the sample dimension
+        self.data = add_sample_dim(to_numpy(data), vars_datamodel)
+
+        # the batch size should be equal to the plateau in the model
+        self.batch_size = list(vars_datamodel.values())[0].shape[0].value if len(vars_datamodel)>0 else 1
+
+        # observed variables in the datamodel
+        obs_datamodel = [k for (k,v) in self.query_vars.items() if k in self.data.keys() and v.is_datamodel]
+
+        # actual data size
+        if len(obs_datamodel)>0:
+            d = self.data[obs_datamodel[0]]
+            if isinstance(d, tf.Tensor):
+                self.data_size = d.get_shape().as_list()[0]
+            else:
+                self.data_size = d.shape[0]
+        else:
+            self.data_size = self.batch_size
+
+        # get batches (even if we have a single one)
+        self.batches = create_batches(self.data, vars_datamodel, self.data_size, self.batch_size)
+
+
+
+
+
+
+
     @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.query_vars, 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.query_vars, 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 +167,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.query_vars[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.query_vars, 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)
+
+
+
+
+