Repository: kazewong/flowMC (language: python)

Issue:
Parsing arguments other than x into log-posterior probability functions
Hi guys, thanks for the really cool code that I'm currently hoping to implement into my own workflow! 

Is it possible to parse other arguments other than x into log-posterior probability functions? This is particularly important for when attempting to sample a distribution which is resultant on comparing observed data to a generated model that is dependent on the sampled parameters.

As an example, the docs specify that a potential target distribution might be:
```
def log_posterior(x):
    return -0.5 * jnp.sum(x ** 2)
```

but it seems impossible to be able to write a posterior function that compares to some observed data. As a very simple example:

```
n_dim = 2

data_x = np.arange(1,10)
m = 2
c = 5
observation = m*data_x + c

def log_posterior(x, observation):
    x_model = np.arange(1,10)
    m = x[0]
    c = x[1]
    y = m*x_model+c
    gaussian = -((y-observation)**2)/(2*np.sqrt(observation)**2)
    return jnp.sum(gaussians)
```

Is this a possible posterior function that just isn't recorded in the docs? I appreciate that it would be hypothetically possible to "generate" or load the observed data within the posterior function but this becomes much less practically feasible when generating the data becomes computationally expensive and it makes sense to do this operation only once.



Discussion / hints:
@bjricketts Thanks for bringing up this issue. The current recommended way to compare with the observation is to "pre-baked" the data into the function instead of loading it as an argument.

In your code, since you have defined `observation` in the global scope, I think your `log_posterior` function will still work without the `observation` in the argument. If you want to dynamically generate the data, we would recommend constructing the `log_posterior` function with another function, something similar to this:

```
def make_posterior_function(data):
    def log_posterior(x):
        return f(x,data)
    return log_posterior
```

The reason behind this seemingly complicated syntax is mainly to avoid triggering recompilation when the data size change. Even though the current code does not directly support what you want to do, I think having a version where we can use syntax like `log_posterior(x,observation)` could be useful. Implementing this will require some thoughts on how to safeguard against unnecessary recompilation, and it will take some work in restructuring the sampler.

There are some community examples out there, here is one of mine on gravitational-wave parameter estimation
https://github.com/kazewong/jim/blob/5f5c51622121cc3b175a0cf0ee00be0f8040b23f/example/ParameterEstimation/GW150914.py#L148
At some point in the future, we are going to link these community examples in the doc.

Please let me know if this is sufficient for your use case. For now, we do not have the bandwidth to restructure the code to accommodate the suggested syntax, but we may come back to this issue at some point.
Thanks @kazewong ! I think this is sufficient for my use case (or at least I can work around the issue with a bit of legwork on my end). This was mostly just a suggestion as there is such a function in the emcee package which enabled me to "generalize" some work I am currently doing with hierarchical Bayesian statistics and reducing the amount of work for me to rewrite or add to some very complex hierarchical Bayesian structures.

While I did only suggest a use case where only observations are an external argument, my own work generally parses a large number of custom functions that fit into another overarching function. Being able to parse series of functions into another function for use is where this use case becomes quite powerful. That being said, I totally understand the fact that this would require a major restructure in the code. I really hope you do get the time at some point to implement it though! 
@bjricketts — Another tip is that you could update your example above as follows:

```python
n_dim = 2

data_x = np.arange(1,10)
m = 2
c = 5
observation = m*data_x + c

def log_posterior(x, observation):
    x_model = np.arange(1,10)
    m = x[0]
    c = x[1]
    y = m*x_model+c
    gaussian = -((y-observation)**2)/(2*jnp.sqrt(observation)**2)
    return jnp.sum(gaussians)

from functools import partial
log_prob_func = partial(log_posterior, observation=observation)
```

(Where all I've done is added the two final lines, and changed a `np` to a `jnp` in your function.)Then use the existing implementation (with `log_prob_func` as the input to the sampler) without any change. This doesn't seem to onerous to me!
That's a nice solution @dfm! That seems like the best course of action for me for now. 

That being said, a perhaps hacky but generalized solution to this issue for implementation into flowMC might be to use something like this:

```
class _FunctionWrapper(object):
    def __init__(self, f, args, kwargs):
        self.f = f
        self.args = args or []
        self.kwargs = kwargs or {}

    def __call__(self, x):
        return self.f(x, *self.args, **self.kwargs)
```

where the logpdf is instead replaced by a _FunctionWrapper object when initalizing the sampler. For instance, in the case of MALA, it could be changed to be:
```
class MALA(LocalSamplerBase):

    def __init__(self, logpdf: Callable, jit: bool, params: dict, args = [], kwargs = {}) -> Callable:
        super().__init__(logpdf, jit, params)
        self.params = params
        self.logpdf = _FunctionWrapper(logpdf,args,kwargs)
```

where args and kwargs are optional arguments that can be parsed through and called as normal. I think this would probably mean that flowMC would need minimal/no rewriting with this solution and users can arbitrarily add as many arguments as they wish to their log posterior functions. Whether this will work with JAX, I'll admit I don't know.

Fix patch:
diff --git a/docs/source/api/flowMC.nfmodel.rst b/docs/source/api/flowMC.nfmodel.rst
index 040a943..aadac82 100644
--- a/docs/source/api/flowMC.nfmodel.rst
+++ b/docs/source/api/flowMC.nfmodel.rst
@@ -4,6 +4,14 @@ flowMC.nfmodel package
 Submodules
 ----------
 
+flowMC.nfmodel.mlp module
+-------------------------
+
+.. automodule:: flowMC.nfmodel.mlp
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 flowMC.nfmodel.realNVP module
 -----------------------------
 
diff --git a/docs/source/api/flowMC.sampler.rst b/docs/source/api/flowMC.sampler.rst
index 9814eda..ef8f1a2 100644
--- a/docs/source/api/flowMC.sampler.rst
+++ b/docs/source/api/flowMC.sampler.rst
@@ -12,6 +12,22 @@ flowMC.sampler.Gaussian\_random\_walk module
    :undoc-members:
    :show-inheritance:
 
+flowMC.sampler.HMC module
+-------------------------
+
+.. automodule:: flowMC.sampler.HMC
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+flowMC.sampler.LocalSampler\_Base module
+----------------------------------------
+
+.. automodule:: flowMC.sampler.LocalSampler_Base
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 flowMC.sampler.MALA module
 --------------------------
 
diff --git a/docs/source/api/flowMC.utils.rst b/docs/source/api/flowMC.utils.rst
index 3a3a0eb..5568f8a 100644
--- a/docs/source/api/flowMC.utils.rst
+++ b/docs/source/api/flowMC.utils.rst
@@ -20,14 +20,6 @@ flowMC.utils.PythonFunctionWrap module
    :undoc-members:
    :show-inheritance:
 
-flowMC.utils.progressBar module
--------------------------------
-
-.. automodule:: flowMC.utils.progressBar
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
 Module contents
 ---------------
 
diff --git a/docs/source/api/modules.rst b/docs/source/api/modules.rst
index cafa3db..3ead41c 100644
--- a/docs/source/api/modules.rst
+++ b/docs/source/api/modules.rst
@@ -1,5 +1,5 @@
-src
-===
+flowMC
+======
 
 .. toctree::
    :maxdepth: 4
diff --git a/docs/source/configuration.rst b/docs/source/configuration.rst
index cd08482..02e8c7a 100644
--- a/docs/source/configuration.rst
+++ b/docs/source/configuration.rst
@@ -15,11 +15,11 @@ This page contains information about the most important hyperparameters which af
 +-----------------------+--------------------------+-----------------------+
 | :ref:`local_sampler`  | :ref:`n_loop_training`   | :ref:`nf_variable`    |
 +-----------------------+--------------------------+-----------------------+
-| :ref:`sampler_params` | :ref:`n_loop_production` | :ref:`local_autotune` |
+| :ref:`data`           | :ref:`n_loop_production` | :ref:`local_autotune` |
 +-----------------------+--------------------------+-----------------------+
-| :ref:`likelihood`     | :ref:`n_local_steps`     | :ref:`train_thinning` |
+| :ref:`nf_model`       | :ref:`n_local_steps`     | :ref:`train_thinning` |
 +-----------------------+--------------------------+-----------------------+
-| :ref:`nf_model`       | :ref:`n_global_steps`    |                       |
+|                       | :ref:`n_global_steps`    |                       |
 +-----------------------+--------------------------+-----------------------+
 |                       | :ref:`n_epochs`          |                       |
 +-----------------------+--------------------------+-----------------------+
@@ -51,6 +51,15 @@ rng_key_set
 
 The set of Jax generated PRNG_keys.
 
+.. _data:
+
+data
+^^^^
+
+The data you want to sample from. This is used to precompile the kernels used during the sampling.
+Note that you keep the shape of the data consistent between runs, otherwise it would trigger recompilation.
+If your likelihood does not take any data arguments, simply put it as None should work.
+
 .. _local_sampler:
 
 local_sampler
@@ -58,20 +67,6 @@ local_sampler
 
 Specific local sampler you want to use.
 
-.. _sampler_params:
-
-sampler_params
-^^^^^^^^^^^^^^
-
-Specific parameters from a particular local sampler.
-
-.. _likelihood:
-
-likelihood
-^^^^^^^^^^
-
-Target log-probability density function you want to sample.
-
 .. _nf_model:
 
 nf_model
diff --git a/docs/source/quickstart.rst b/docs/source/quickstart.rst
index c5ccbfa..6b6354c 100644
--- a/docs/source/quickstart.rst
+++ b/docs/source/quickstart.rst
@@ -49,8 +49,10 @@ To sample a N dimensional Gaussian, you would do something like:
     from flowMC.utils.PRNG_keys import initialize_rng_keys
     from flowMC.nfmodel.utils import *
 
-    def log_posterior(x):
-        return -0.5 * jnp.sum(x ** 2)
+    def log_posterior(x, data):
+        return -0.5 * jnp.sum((x-data) ** 2)
+
+    data = jnp.arange(5)
 
     n_dim = 5
     n_chains = 10
@@ -63,8 +65,8 @@ To sample a N dimensional Gaussian, you would do something like:
 
     nf_sampler = Sampler(n_dim,
                         rng_key_set,
+                        jnp.arange(n_dim),
                         local_sampler,
-                        log_posterior,
                         model,
                         n_local_steps = 50,
                         n_global_steps = 50,
@@ -73,12 +75,11 @@ To sample a N dimensional Gaussian, you would do something like:
                         batch_size = 1000,
                         n_chains = n_chains)
 
-    nf_sampler.sample(initial_position)
+    nf_sampler.sample(initial_position, data)
     chains,log_prob,local_accs, global_accs = nf_sampler.get_sampler_state().values()
 
-For more examples, have a look at the :ref:`tutorials`. 
-.. `GitHub <https://github.com/kazewong/flowMC/tree/main/example>`_.
-.. In particular, currently the best engineered test case is `dualmoon.py <https://github.com/kazewong/flowMC/blob/main/example/dualmoon.py>`_.
+For more examples, have a look at the :ref:`tutorials` on `GitHub <https://github.com/kazewong/flowMC/tree/main/example>`_.
+In particular, currently the best engineered test case is `dualmoon.py <https://github.com/kazewong/flowMC/blob/main/example/dualmoon.py>`_.
 
 In the ideal case, the only three things you will have to do are:
 
diff --git a/example/dualmoon.py b/example/dualmoon.py
index 468cb82..82004e6 100644
--- a/example/dualmoon.py
+++ b/example/dualmoon.py
@@ -5,7 +5,6 @@
 from jax.scipy.special import logsumexp
 import numpy as np
 
-from flowMC.nfmodel.realNVP import RealNVP
 from flowMC.nfmodel.rqSpline import RQSpline
 from flowMC.nfmodel.utils import *
 from flowMC.sampler.MALA import MALA
@@ -13,17 +12,16 @@
 from flowMC.utils.PRNG_keys import initialize_rng_keys
 
 
-def dual_moon_pe(x):
+def dual_moon_pe(x, data):
     """
     Term 2 and 3 separate the distribution and smear it along the first and second dimension
     """
-    term1 = 0.5 * ((jnp.linalg.norm(x) - 2) / 0.1) ** 2
+    print("compile count")
+    term1 = 0.5 * ((jnp.linalg.norm(x-data) - 2) / 0.1) ** 2
     term2 = -0.5 * ((x[:1] + jnp.array([-3.0, 3.0])) / 0.8) ** 2
     term3 = -0.5 * ((x[1:2] + jnp.array([-3.0, 3.0])) / 0.6) ** 2
     return -(term1 - logsumexp(term2) - logsumexp(term3))
 
-### Demo config
-
 n_dim = 5
 n_chains = 20
 n_loop_training = 5
@@ -35,25 +33,22 @@ def dual_moon_pe(x):
 num_epochs = 30
 batch_size = 10000
 
-print("Preparing RNG keys")
-rng_key_set = initialize_rng_keys(n_chains, seed=42)
+data = jnp.zeros(n_dim)
 
-print("Initializing MCMC model and normalizing flow model.")
+rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
 initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
 
-model = RQSpline(n_dim, 10, [128, 128], 8)
-
-
-MALA_Sampler = MALA(dual_moon_pe, True, {"step_size": 1e-1})
+MALA_Sampler = MALA(dual_moon_pe, True, {"step_size": 0.1})
+model = RQSpline(n_dim, 4, [32, 32], 8)
 
 print("Initializing sampler class")
 
 nf_sampler = Sampler(
     n_dim,
     rng_key_set,
+    jnp.zeros(5),
     MALA_Sampler,
-    dual_moon_pe,
     model,
     n_loop_training=n_loop_training,
     n_loop_production=n_loop_production,
@@ -66,10 +61,8 @@ def dual_moon_pe(x):
     batch_size=batch_size,
     use_global=True,
 )
-print("Sampling")
-
-nf_sampler.sample(initial_position)
 
+nf_sampler.sample(initial_position, data)
 summary = nf_sampler.get_sampler_state(training=True)
 chains, log_prob, local_accs, global_accs, loss_vals = summary.values() 
 nf_samples = nf_sampler.sample_flow(10000)
diff --git a/src/flowMC/sampler/Gaussian_random_walk.py b/src/flowMC/sampler/Gaussian_random_walk.py
index 587eb64..f0ab315 100644
--- a/src/flowMC/sampler/Gaussian_random_walk.py
+++ b/src/flowMC/sampler/Gaussian_random_walk.py
@@ -19,17 +19,23 @@ def __init__(self, logpdf: Callable, jit: bool, params: dict, verbose: bool = Fa
         super().__init__(logpdf, jit, params)
         self.params = params
         self.logpdf = logpdf
+        self.logpdf_vmap = jax.vmap(logpdf, in_axes=(0, None))
         self.verbose = verbose
+        self.kernel = None
+        self.kernel_vmap = None
+        self.update = None
+        self.update_vmap = None
+        self.sampler = None
 
     def make_kernel(self, return_aux = False) -> Callable:
         """
         Making a single update of the random walk
         """
-        def rw_kernel(rng_key, position, log_prob, params = {"step_size": 0.1}):
+        def rw_kernel(rng_key, position, log_prob, data, params = {"step_size": 0.1}):
             key1, key2 = jax.random.split(rng_key)
             move_proposal = jax.random.normal(key1, shape=position.shape) * params['step_size']
             proposal = position + move_proposal
-            proposal_log_prob = self.logpdf(proposal)
+            proposal_log_prob = self.logpdf(proposal, data)
 
             log_uniform = jnp.log(jax.random.uniform(key2))
             do_accept = log_uniform < proposal_log_prob - log_prob
@@ -46,16 +52,17 @@ def make_update(self) -> Callable:
         Making a the random walk update function for multiple steps
         """
 
-        rw_kernel = self.make_kernel()
+        if self.kernel is None:
+            raise ValueError("Kernel not defined. Please run make_kernel first.")
 
         def rw_update(i, state):
-            key, positions, log_p, acceptance, params = state
+            key, positions, log_p, acceptance, data, params = state
             _, key = jax.random.split(key)
-            new_position, new_log_p, do_accept = rw_kernel(key, positions[i-1], log_p[i-1], params)
+            new_position, new_log_p, do_accept = self.kernel(key, positions[i-1], log_p[i-1], data, params)
             positions = positions.at[i].set(new_position)
             log_p = log_p.at[i].set(new_log_p)
             acceptance = acceptance.at[i].set(do_accept)
-            return (key, positions, log_p, acceptance, params)
+            return (key, positions, log_p, acceptance, data, params)
         
         return rw_update
 
@@ -63,31 +70,25 @@ def make_sampler(self) -> Callable:
         """
         Making the random walk sampler for multiple chains given initial positions
         """
-        rw_update = self.make_update()
-        lp = self.logpdf
+        if self.update is None:
+            raise ValueError("Update function not defined. Please run make_update first.")
 
-        if self.jit:
-            rw_update = jax.jit(rw_update)
-            lp = jax.jit(self.logpdf)
 
-        rw_update = jax.vmap(rw_update, in_axes = (None, (0, 0, 0, 0, None)), out_axes=(0, 0, 0, 0, None))
-        lp = jax.vmap(lp)
-
-        def rw_sampler(rng_key, n_steps, initial_position):
-            logp = lp(initial_position)
+        def rw_sampler(rng_key, n_steps, initial_position, data):
+            logp = self.logpdf_vmap(initial_position, data)
             n_chains = rng_key.shape[0]
             acceptance = jnp.zeros((n_chains, n_steps))
             all_positions = (jnp.zeros((n_chains, n_steps) + initial_position.shape[-1:])) + initial_position[:, None]
             all_logp = (jnp.zeros((n_chains, n_steps)) + logp[:, None])
-            state = (rng_key, all_positions, all_logp, acceptance, self.params)
+            state = (rng_key, all_positions, all_logp, acceptance, data, self.params)
             if self.verbose:
                 iterator_loop = tqdm(range(1, n_steps), desc="Sampling Locally", miniters=int(n_steps / 10))
             else:
                 iterator_loop = range(1, n_steps)
 
             for i in iterator_loop:
-                state = rw_update(i, state)
-            return state[:-1]
+                state = self.update_vmap(i, state)
+            return state[:-2]
             
 
         return rw_sampler
diff --git a/src/flowMC/sampler/HMC.py b/src/flowMC/sampler/HMC.py
index 170f6b7..62b40ec 100644
--- a/src/flowMC/sampler/HMC.py
+++ b/src/flowMC/sampler/HMC.py
@@ -19,7 +19,7 @@ class HMC(LocalSamplerBase):
     def __init__(self, logpdf: Callable, jit: bool, params: dict, verbose: bool = False) -> Callable:
         super().__init__(logpdf, jit, params)
 
-        self.potential = lambda x: -self.logpdf(x)
+        self.potential = lambda x, data: -logpdf(x, data)
         self.grad_potential = jax.grad(self.potential)
         
         self.params = params
@@ -38,9 +38,17 @@ def __init__(self, logpdf: Callable, jit: bool, params: dict, verbose: bool = Fa
 
         self.kinetic = lambda p, params: 0.5*(p**2 * params['inverse_metric']).sum()
         self.grad_kinetic = jax.grad(self.kinetic)
+        self.logpdf = self.potential
+        self.logpdf_vmap = jax.vmap(self.logpdf, in_axes=(0, None))
         self.verbose = verbose
+        self.kernel = None
+        self.kernel_vmap = None
+        self.update = None
+        self.update_vmap = None
+        self.sampler = None
 
     def get_initial_hamiltonian(self, rng_key: jax.random.PRNGKey, position: jnp.array,
+                                data: jnp.array,
                                 params: dict):
         """
         Compute the value of the Hamiltonian from positions with initial momentum draw
@@ -48,28 +56,28 @@ def get_initial_hamiltonian(self, rng_key: jax.random.PRNGKey, position: jnp.arr
         """
         
         momentum = jax.random.normal(rng_key, shape=position.shape) * params['inverse_metric'] **-0.5
-        return self.potential(position) + self.kinetic(momentum, params)
+        return self.potential(position, data) + self.kinetic(momentum, params)
 
     def make_kernel(self, return_aux = False) -> Callable:
         """
         Making HMC kernel for a single step
         """
 
-        def leapfrog_kernel(carry, data):
-            position, momentum, params = carry
+        def leapfrog_kernel(carry, extras):
+            position, momentum, params, data = carry
             position = position + params['step_size'] * self.grad_kinetic(momentum, params)
-            momentum = momentum - params['step_size'] * self.grad_potential(position)
-            return (position, momentum, params), data
+            momentum = momentum - params['step_size'] * self.grad_potential(position, data)
+            return (position, momentum, params, data), extras
 
 
-        def leapfrog_step(position, momentum, params: dict):
-            momentum = momentum - 0.5 * params['step_size'] * self.grad_potential(position)
-            (position, momentum, params), _ = jax.lax.scan(leapfrog_kernel, (position, momentum, params), jnp.arange(self.n_leapfrog-1))
+        def leapfrog_step(position, momentum, data, params: dict):
+            momentum = momentum - 0.5 * params['step_size'] * self.grad_potential(position, data)
+            (position, momentum, params, data), _ = jax.lax.scan(leapfrog_kernel, (position, momentum, params, data), jnp.arange(self.n_leapfrog-1))
             position = position + params['step_size'] * self.grad_kinetic(momentum, params)
-            momentum = momentum - 0.5*params['step_size'] * self.grad_potential(position)
+            momentum = momentum - 0.5*params['step_size'] * self.grad_potential(position, data)
             return position, momentum
 
-        def hmc_kernel(rng_key, position, PE, params: dict):
+        def hmc_kernel(rng_key, position, PE, data, params: dict):
             """
             Note that since the potential function is the negative log likelihood, 
             hamiltonian is going down, but the likelihood value should go up.
@@ -83,8 +91,8 @@ def hmc_kernel(rng_key, position, PE, params: dict):
 
             momentum = jax.random.normal(key1, shape=position.shape) * params['inverse_metric']**-0.5
             H = PE + self.kinetic(momentum, params)
-            proposed_position, proposed_momentum = leapfrog_step(position, momentum, params)
-            proposed_PE = self.potential(proposed_position)
+            proposed_position, proposed_momentum = leapfrog_step(position, momentum, data, params)
+            proposed_PE = self.potential(proposed_position, data)
             proposed_ham = proposed_PE + self.kinetic(proposed_momentum, params)
             log_acc = H - proposed_ham
             log_uniform = jnp.log(jax.random.uniform(key2))
@@ -107,18 +115,19 @@ def make_update(self) -> Callable:
         Making HMC update function for multiple steps
         """
 
-        hmc_kernel = self.make_kernel()
+        if self.kernel is None:
+            raise ValueError("Kernel not defined. Please run make_kernel first.")
 
         def hmc_update(i, state):
-            key, positions, PE, acceptance, params = state
+            key, positions, PE, acceptance, data, params = state
             _, key = jax.random.split(key)
-            new_position, new_PE, do_accept = hmc_kernel(
-                key, positions[i - 1], PE[i - 1], params
+            new_position, new_PE, do_accept = self.kernel(
+                key, positions[i - 1], PE[i - 1], data, params
             )
             positions = positions.at[i].set(new_position)
             PE = PE.at[i].set(new_PE)
             acceptance = acceptance.at[i].set(do_accept)
-            return (key, positions, PE, acceptance, params)
+            return (key, positions, PE, acceptance, data, params)
 
         return hmc_update
 
@@ -127,20 +136,15 @@ def make_sampler(self) -> Callable:
         Making HMC sampler function for multiple chains from initial positions
         """
 
-        hmc_update = self.make_update()
+        if self.update is None:
+            raise ValueError("Update function not defined. Please run make_update first.")
 
-        if self.jit:
-            hmc_update = jax.jit(hmc_update)
-
-        hmc_update = jax.vmap(hmc_update, in_axes = (None, (0, 0, 0, 0, None)), out_axes=(0, 0, 0, 0, None))
-        lp = jax.vmap(self.potential)
-
-        def hmc_sampler(rng_key, n_steps, initial_position):
+        def hmc_sampler(rng_key, n_steps, initial_position, data):
             
             keys = jax.vmap(jax.random.split)(rng_key)
             rng_key = keys[:, 0]
             rng_init = keys[:, 1]
-            logp = lp(initial_position)
+            logp = self.logpdf_vmap(initial_position, data)
             n_chains = rng_key.shape[0]
             acceptance = jnp.zeros(
                 (
@@ -167,7 +171,7 @@ def hmc_sampler(rng_key, n_steps, initial_position):
                 )
                 + logp[:, None]
             )
-            state = (rng_key, all_positions, all_logp, acceptance, self.params)
+            state = (rng_key, all_positions, all_logp, acceptance, data, self.params)
 
             if self.verbose:
                 iterator_loop = tqdm(range(1, n_steps), desc="Sampling Locally", miniters=int(n_steps / 10))
@@ -175,7 +179,7 @@ def hmc_sampler(rng_key, n_steps, initial_position):
                 iterator_loop = range(1, n_steps)
 
             for i in iterator_loop:
-                state = hmc_update(i, state)
+                state = self.update_vmap(i, state)
 
             state = (state[0], state[1], -state[2], state[3])
             return state
@@ -183,9 +187,9 @@ def hmc_sampler(rng_key, n_steps, initial_position):
         return hmc_sampler
 
     def make_leapforg_kernel(self):
-        def leapfrog_kernel(carry, data):
-            position, momentum, params = carry
+        def leapfrog_kernel(carry, extras):
+            position, momentum, params, data = carry
             position = position + self.step_size * self.grad_kinetic(momentum, params)
-            momentum = momentum - self.step_size * self.grad_potential(position)
+            momentum = momentum - self.step_size * self.grad_potential(position, data)
             return position, momentum
         return leapfrog_kernel
diff --git a/src/flowMC/sampler/LocalSampler_Base.py b/src/flowMC/sampler/LocalSampler_Base.py
index 62b95e2..757a3ca 100644
--- a/src/flowMC/sampler/LocalSampler_Base.py
+++ b/src/flowMC/sampler/LocalSampler_Base.py
@@ -1,5 +1,7 @@
 from abc import abstractmethod
 from typing import Callable
+import jax
+import jax.numpy as jnp
 class LocalSamplerBase:
 
     def __init__(self, logpdf: Callable, jit: bool, params: dict) -> Callable:
@@ -10,6 +12,35 @@ def __init__(self, logpdf: Callable, jit: bool, params: dict) -> Callable:
         self.jit = jit
         self.params = params
 
+    def precompilation(self, n_chains, n_dims, n_step, data):
+
+        if self.jit == True:
+            print("jit is requested, precompiling kernels and update...")
+        else:
+            print("jit is not requested, compiling only vmap functions...")
+
+        self.kernel = self.make_kernel()
+        self.kernel_vmap = jax.vmap(self.kernel, in_axes = (0, 0, 0, None, None), out_axes=(0, 0, 0))
+        self.update = self.make_update()
+        self.update_vmap = jax.vmap(self.update, in_axes = (None, (0, 0, 0, 0, None, None)), out_axes=(0, 0, 0, 0, None, None))
+        self.sampler = self.make_sampler()
+
+        if self.jit == True:
+            self.logpdf_vmap = jax.jit(self.logpdf_vmap)
+            self.kernel = jax.jit(self.kernel)
+            self.kernel_vmap = jax.jit(self.kernel_vmap)
+            self.update = jax.jit(self.update)
+            self.update_vmap = jax.jit(self.update_vmap)
+            self.kernel(jax.random.PRNGKey(0), jnp.ones(n_dims), jnp.ones(1), data, self.params)
+            # self.update(1, (jax.random.PRNGKey(0), jnp.ones(n_dims), jnp.ones(1), jnp.zeros((n_step, 1)), data, self.params))
+        
+        key = jax.random.split(jax.random.PRNGKey(0), n_chains)
+
+        self.logpdf_vmap(jnp.ones((n_chains, n_dims)), data)
+        self.kernel_vmap(key, jnp.ones((n_chains, n_dims)), jnp.ones((n_chains, 1)), data, self.params)
+        self.update_vmap(1, (key, jnp.ones((n_chains, n_step, n_dims)), jnp.ones((n_chains, n_step, 1)),jnp.zeros((n_chains, n_step, 1)), data, self.params))
+        
+
     @abstractmethod
     def make_kernel(self, return_aux = False) -> Callable:
         """
diff --git a/src/flowMC/sampler/MALA.py b/src/flowMC/sampler/MALA.py
index 75b2640..9b335cf 100644
--- a/src/flowMC/sampler/MALA.py
+++ b/src/flowMC/sampler/MALA.py
@@ -4,6 +4,7 @@
 from jax.scipy.stats import multivariate_normal
 from tqdm import tqdm
 from flowMC.sampler.LocalSampler_Base import LocalSamplerBase
+from functools import partialmethod
 
 
 class MALA(LocalSamplerBase):
@@ -16,11 +17,18 @@ class MALA(LocalSamplerBase):
         params: dictionary of parameters for the sampler
     """
 
-    def __init__(self, logpdf: Callable, jit: bool, params: dict, verbose: bool = False) -> Callable:
+    def __init__(self, logpdf: Callable, jit: bool, params: dict, verbose: bool = False, use_autotune = False) -> Callable:
         super().__init__(logpdf, jit, params)
         self.params = params
         self.logpdf = logpdf
+        self.logpdf_vmap = jax.vmap(logpdf, in_axes=(0, None))
         self.verbose = verbose
+        self.kernel = None
+        self.kernel_vmap = None
+        self.update = None
+        self.update_vmap = None
+        self.sampler = None
+        self.use_autotune = use_autotune
 
     def make_kernel(self, return_aux = False) -> Callable:
         """
@@ -33,14 +41,14 @@ def make_kernel(self, return_aux = False) -> Callable:
             mala_kernel (Callable) A MALA kernel.
         """
         def body(carry, this_key):
-            this_position, dt = carry
+            this_position, dt, data = carry
             dt2 = dt*dt
-            this_log_prob, this_d_log = jax.value_and_grad(self.logpdf)(this_position)
+            this_log_prob, this_d_log = jax.value_and_grad(self.logpdf)(this_position, data)
             proposal = this_position + jnp.dot(dt2, this_d_log) / 2
             proposal += jnp.dot(dt, jax.random.normal(this_key, shape=this_position.shape))
-            return (proposal,dt), (proposal, this_log_prob, this_d_log)
+            return (proposal,dt, data), (proposal, this_log_prob, this_d_log)
 
-        def mala_kernel(rng_key, position, log_prob, params = {"step_size": 0.1}):
+        def mala_kernel(rng_key, position, log_prob, data, params = {"step_size": 0.1}):
             """
             Metropolis-adjusted Langevin algorithm kernel.
             This function make a proposal and accept/reject it.
@@ -49,6 +57,8 @@ def mala_kernel(rng_key, position, log_prob, params = {"step_size": 0.1}):
                 rng_key (n_chains, 2): random key
                 position (n_chains, n_dim): current position
                 log_prob (n_chains, ): log-probability of the current position
+                data: data to be passed to the logpdf
+                params: dictionary of parameters for the sampler
 
             Returns:
                 position (n_chains, n_dim): the new poisiton of the chain
@@ -62,7 +72,7 @@ def mala_kernel(rng_key, position, log_prob, params = {"step_size": 0.1}):
             dt2 = dt * dt
 
             _, (proposal, logprob, d_logprob) = jax.lax.scan(
-                body, (position, dt), jnp.array([key1, key1])
+                body, (position, dt, data), jnp.array([key1, key1])
             )
 
             ratio = logprob[1] - logprob[0]
@@ -86,16 +96,17 @@ def make_update(self) -> Callable:
         """
         Make a MALA update function for multiple steps
         """
-        mala_kernel = self.make_kernel()
+        if self.kernel is None:
+            raise ValueError("Kernel not defined. Please run make_kernel first.")
 
         def mala_update(i, state):
-            key, positions, log_p, acceptance, params = state
+            key, positions, log_p, acceptance, data, params = state
             _, key = jax.random.split(key)
-            new_position, new_log_p, do_accept = mala_kernel(key, positions[i-1], log_p[i-1], params)
+            new_position, new_log_p, do_accept = self.kernel(key, positions[i-1], log_p[i-1], data, params)
             positions = positions.at[i].set(new_position)
             log_p = log_p.at[i].set(new_log_p)
             acceptance = acceptance.at[i].set(do_accept)
-            return (key, positions, log_p, acceptance, params)
+            return (key, positions, log_p, acceptance, data, params)
         
         return mala_update
 
@@ -103,64 +114,57 @@ def make_sampler(self) -> Callable:
         """
         Make a MALA sampler for multiple chains given initial positions
         """
-        mala_update = self.make_update()
-        lp = self.logpdf
-
-        if self.jit:
-            mala_update = jax.jit(mala_update)
-            lp = jax.jit(self.logpdf)
 
-        mala_update = jax.vmap(mala_update, in_axes = (None, (0, 0, 0, 0, None)), out_axes=(0, 0, 0, 0, None))
-        lp = jax.vmap(lp)
+        if self.update is None:
+            raise ValueError("Update function not defined. Please run make_update first.")
 
-        def mala_sampler(rng_key, n_steps, initial_position):
-            logp = lp(initial_position)
+        def mala_sampler(rng_key, n_steps, initial_position, data):
+            logp = self.logpdf_vmap(initial_position, data)
             n_chains = rng_key.shape[0]
             acceptance = jnp.zeros((n_chains, n_steps))
             all_positions = (jnp.zeros((n_chains, n_steps) + initial_position.shape[-1:])) + initial_position[:, None]
             all_logp = (jnp.zeros((n_chains, n_steps)) + logp[:, None])
-            state = (rng_key, all_positions, all_logp, acceptance, self.params)
+            state = (rng_key, all_positions, all_logp, acceptance, data, self.params)
             if self.verbose:
                 iterator_loop = tqdm(range(1, n_steps), desc="Sampling Locally", miniters=int(n_steps / 10))
             else:
                 iterator_loop = range(1, n_steps)
             for i in iterator_loop:
-                state = mala_update(i, state)
-            return state[:-1]
+                state = self.update_vmap(i, state)
+            return state[:-2]
 
+        self.sampler = mala_sampler
         return mala_sampler 
 
-from tqdm import tqdm
-from functools import partialmethod
 
-def mala_sampler_autotune(mala_kernel_vmap, rng_key, initial_position, log_prob, params, max_iter = 30):
-    """
-    Tune the step size of the MALA kernel using the acceptance rate.
+    def mala_sampler_autotune(self, rng_key, initial_position, log_prob, data, params, max_iter = 30):
+        """
+        Tune the step size of the MALA kernel using the acceptance rate.
 
-    Args:
-        mala_kernel_vmap (Callable): A MALA kernel
-        rng_key: Jax PRNGKey
-        initial_position (n_chains, n_dim): initial position of the chains
-        log_prob (n_chains, ): log-probability of the initial position
-        params (dict): parameters of the MALA kernel
-        max_iter (int): maximal number of iterations to tune the step size
-    """
+        Args:
+            mala_kernel_vmap (Callable): A MALA kernel
+            rng_key: Jax PRNGKey
+            initial_position (n_chains, n_dim): initial position of the chains
+            log_prob (n_chains, ): log-probability of the initial position
+            params (dict): parameters of the MALA kernel
+            max_iter (int): maximal number of iterations to tune the step size
+        """
+
+        tqdm.__init__ = partialmethod(tqdm.__init__, disable=True)
 
-    tqdm.__init__ = partialmethod(tqdm.__init__, disable=True)
-
-    counter = 0
-    position, log_prob, do_accept = mala_kernel_vmap(rng_key, initial_position, log_prob, params)
-    acceptance_rate = jnp.mean(do_accept)
-    while (acceptance_rate <= 0.3) or (acceptance_rate >= 0.5):
-        if counter > max_iter:
-            print("Maximal number of iterations reached. Existing tuning with current parameters.")
-            break
-        if acceptance_rate <= 0.3:
-            params['step_size'] *= 0.8
-        elif acceptance_rate >= 0.5:
-            params['step_size'] *= 1.25
-        counter += 1
-        position, log_prob, do_accept = mala_kernel_vmap(rng_key, initial_position, log_prob, params)
+        counter = 0
+        position, log_prob, do_accept = self.kernel_vmap(rng_key, initial_position, log_prob, data, params)
         acceptance_rate = jnp.mean(do_accept)
-    tqdm.__init__ = partialmethod(tqdm.__init__, disable=False)
-    return params
\ No newline at end of file
+        while (acceptance_rate <= 0.3) or (acceptance_rate >= 0.5):
+            if counter > max_iter:
+                print("Maximal number of iterations reached. Existing tuning with current parameters.")
+                break
+            if acceptance_rate <= 0.3:
+                params['step_size'] *= 0.8
+            elif acceptance_rate >= 0.5:
+                params['step_size'] *= 1.25
+            counter += 1
+            position, log_prob, do_accept = self.kernel_vmap(rng_key, initial_position, log_prob, data, params)
+            acceptance_rate = jnp.mean(do_accept)
+        tqdm.__init__ = partialmethod(tqdm.__init__, disable=False)
+        return params
diff --git a/src/flowMC/sampler/NF_proposal.py b/src/flowMC/sampler/NF_proposal.py
index 23f140e..bb0ef0d 100644
--- a/src/flowMC/sampler/NF_proposal.py
+++ b/src/flowMC/sampler/NF_proposal.py
@@ -117,6 +117,7 @@ def make_nf_metropolis_sampler(nf_model):
                 nf_variables: Normalizing flow variables.
                 target_pdf: Target pdf.
                 initial_position: Initial position.
+                data: Data.
 
             Returns:
                 rng_key: Current state of random key
@@ -146,7 +147,7 @@ def sample_nf(rng_key, n_samples, nf_params, nf_variables):
     sample_nf = jax.jit(sample_nf, static_argnums=(1))
 
     def nf_metropolis_sampler(
-        rng_key, n_steps, nf_param, nf_variables, target_pdf, initial_position
+        rng_key, n_steps, nf_param, nf_variables, target_pdf, initial_position, data
     ):
         rng_key, *subkeys = random.split(rng_key, 3)
 
@@ -155,7 +156,7 @@ def nf_metropolis_sampler(
         log_pdf_nf_initial = eval_nf_logprob(
             initial_position, nf_param, nf_variables
         )
-        log_pdf_initial = target_pdf(initial_position)
+        log_pdf_initial = target_pdf(initial_position, data)
 
         if total_sample > n_sample_max:
             proposal_position = jnp.zeros((total_sample, initial_position.shape[-1]))
@@ -176,7 +177,7 @@ def nf_metropolis_sampler(
                 ].set(eval_nf_logprob(local_samples, nf_param, nf_variables))
                 log_pdf_proposal = log_pdf_proposal.at[
                     i * n_sample_max : (i + 1) * n_sample_max
-                ].set(target_pdf(local_samples))
+                ].set(target_pdf(local_samples, data))
                 local_key, subkey = random.split(local_key, 2)
 
         else:
@@ -186,7 +187,7 @@ def nf_metropolis_sampler(
             log_pdf_nf_proposal = eval_nf_logprob(
                 proposal_position, nf_param, nf_variables
             )
-            log_pdf_proposal = target_pdf(proposal_position)
+            log_pdf_proposal = target_pdf(proposal_position, data)
 
         proposal_position = proposal_position.reshape(
             n_steps, initial_position.shape[0], initial_position.shape[1]
diff --git a/src/flowMC/sampler/Sampler.py b/src/flowMC/sampler/Sampler.py
index 77f5346..1f08060 100644
--- a/src/flowMC/sampler/Sampler.py
+++ b/src/flowMC/sampler/Sampler.py
@@ -39,15 +39,15 @@ class Sampler():
         keep_quantile (float, optional): Quantile of chains to keep when training the normalizing flow model. Defaults to 0..
         local_autotune (None, optional): Auto-tune function for the local sampler. Defaults to None.
         train_thinning (int, optional): Thinning for the data used to train the normalizing flow. Defaults to 1.
-        model_init (dic, optional): Dictionnary with keys "params" and "variables" for the NF model. Defaults to None.
+        model_init (dic, optional): Dictionary with keys "params" and "variables" for the NF model. Defaults to None.
     """
 
     def __init__(
         self,
         n_dim: int,
         rng_key_set: Tuple,
+        data: jnp.ndarray,
         local_sampler: LocalSamplerBase,
-        likelihood: Callable,
         nf_model: Callable,
         n_loop_training: int = 3,
         n_loop_production: int = 3,
@@ -69,11 +69,7 @@ def __init__(
     ):
         rng_key_init, rng_keys_mcmc, rng_keys_nf, init_rng_keys_nf = rng_key_set
 
-        self.likelihood = likelihood
-        self.likelihood_vec = jax.jit(jax.vmap(self.likelihood))
-        self.local_sampler_class = local_sampler
-        self.local_sampler = local_sampler.make_sampler()
-        self.local_autotune = local_autotune
+        # Copying input into the model
 
         self.rng_keys_nf = rng_keys_nf
         self.rng_keys_mcmc = rng_keys_mcmc
@@ -90,8 +86,17 @@ def __init__(
         self.batch_size = batch_size
         self.use_global = use_global
         self.logging = logging
+        self.keep_quantile = keep_quantile
+        self.train_thinning = train_thinning
+
+        # Initialized local and global samplers
 
+        self.local_sampler_class = local_sampler
+        self.local_sampler_class.precompilation(n_chains=n_chains, n_dims=n_dim, n_step=n_local_steps, data=data)
+        self.local_sampler = self.local_sampler_class.sampler
 
+        self.local_autotune = local_autotune
+        self.likelihood_vec = self.local_sampler_class.logpdf_vmap
         self.nf_model = nf_model
         if model_init is None:
             model_init = nf_model.init(init_rng_keys_nf, jnp.ones((1, self.n_dim)))
@@ -99,23 +104,17 @@ def __init__(
         self.variables = model_init["variables"]
         if nf_variable is not None:
             self.variables = self.variables
-
-        self.keep_quantile = keep_quantile
         self.global_sampler = make_nf_metropolis_sampler(self.nf_model)
         
-        self.train_thinning = train_thinning
         self.nf_training_loop, train_epoch, train_step = make_training_loop(
             self.nf_model
         )
         tx = optax.adam(self.learning_rate, self.momentum)
-        # tx = optax.chain(
-        #     optax.adaptive_grad_clip(1, eps=0.001),
-        #     optax.adam(self.learning_rate, self.momentum)
-        # )
         self.state = train_state.TrainState.create(
             apply_fn=nf_model.apply, params=params, tx=tx
         )
 
+        # Initialized result dictionary
         training = {}
         training["chains"] = jnp.empty((self.n_chains, 0, self.n_dim))
         training["log_prob"] = jnp.empty((self.n_chains, 0))
@@ -133,7 +132,7 @@ def __init__(
         self.summary['training'] = training
         self.summary['production'] = production
 
-    def sample(self, initial_position):
+    def sample(self, initial_position, data):
         """
         Sample from the posterior using the local sampler.
 
@@ -151,14 +150,14 @@ def sample(self, initial_position):
         # Note that auto-tune function needs to have the same number of steps
         # as the actual sampling loop to avoid recompilation.
 
-        self.local_sampler_tuning(initial_position)
+        self.local_sampler_tuning(initial_position, data)
         last_step = initial_position
         if self.use_global == True:
-            last_step = self.global_sampler_tuning(last_step)
+            last_step = self.global_sampler_tuning(last_step, data)
 
-        last_step = self.production_run(last_step)
+        last_step = self.production_run(last_step, data)
 
-    def sampling_loop(self, initial_position: jnp.array, training=False) -> jnp.array:
+    def sampling_loop(self, initial_position: jnp.array, data: jnp.array, training=False) -> jnp.array:
         """
         One sampling loop that iterate through the local sampler and potentially the global sampler.
         If training is set to True, the loop will also train the normalizing flow model.
@@ -177,7 +176,7 @@ def sampling_loop(self, initial_position: jnp.array, training=False) -> jnp.arra
             summary_mode = 'production'
 
         self.rng_keys_mcmc, positions, log_prob, local_acceptance = self.local_sampler(
-            self.rng_keys_mcmc, self.n_local_steps, initial_position
+            self.rng_keys_mcmc, self.n_local_steps, initial_position, data
         )
 
         self.summary[summary_mode]['chains'] = jnp.append(
@@ -245,6 +244,7 @@ def sampling_loop(self, initial_position: jnp.array, training=False) -> jnp.arra
                 self.variables,
                 self.likelihood_vec,
                 positions[:, -1],
+                data,
             )
 
             self.summary[summary_mode]['chains'] = jnp.append(
@@ -262,7 +262,7 @@ def sampling_loop(self, initial_position: jnp.array, training=False) -> jnp.arra
 
         return last_step
 
-    def local_sampler_tuning(self, initial_position: jnp.array, max_iter: int = 100):
+    def local_sampler_tuning(self, initial_position: jnp.array, data: jnp.array, max_iter: int = 100):
         """
         Tuning the local sampler. This runs a number of iterations of the local sampler,
         and then uses the acceptance rate to adjust the local sampler parameters.
@@ -276,15 +276,15 @@ def local_sampler_tuning(self, initial_position: jnp.array, max_iter: int = 100)
         """
         if self.local_autotune is not None:
             print("Autotune found, start tuning sampler_params")
-            kernel_vmap = jax.vmap(self.local_sampler_class.make_kernel(), in_axes = (0, 0, 0,  None), out_axes=(0, 0, 0))
-            self.local_sampler_class.params = self.local_autotune(
-                kernel_vmap, self.rng_keys_mcmc, initial_position, self.likelihood_vec(initial_position), self.local_sampler_class.params, max_iter)
-            self.local_sampler = self.local_sampler_class.make_sampler()
+            kernel_vmap = self.local_sampler.kernel_vmap
+            self.local_sampler.params = self.local_autotune(
+                kernel_vmap, self.rng_keys_mcmc, initial_position, self.likelihood_vec(initial_position), data, self.local_sampler.params, max_iter)
+            self.local_sampler = self.local_sampler.make_sampler()
 
         else:
             print("No autotune found, use input sampler_params")
 
-    def global_sampler_tuning(self, initial_position: jnp.ndarray) -> jnp.array:
+    def global_sampler_tuning(self, initial_position: jnp.ndarray, data: jnp.array) -> jnp.array:
         """
         Tuning the global sampler. This runs both the local sampler and the global sampler,
         and train the normalizing flow on the run.
@@ -302,10 +302,10 @@ def global_sampler_tuning(self, initial_position: jnp.ndarray) -> jnp.array:
                 range(self.n_loop_training),
                 desc="Tuning global sampler",
                 ):
-            last_step = self.sampling_loop(last_step, training=True)
+            last_step = self.sampling_loop(last_step, data, training=True)
         return last_step
 
-    def production_run(self, initial_position: jnp.ndarray) -> jnp.array:
+    def production_run(self, initial_position: jnp.ndarray, data: jnp.array) -> jnp.array:
         """
         Sampling procedure that produce the final set of samples.
         The main difference between this and the global tuning step is
@@ -322,7 +322,7 @@ def production_run(self, initial_position: jnp.ndarray) -> jnp.array:
                 range(self.n_loop_production),
                 desc="Production run",
                 ):
-            last_step = self.sampling_loop(last_step)
+            last_step = self.sampling_loop(last_step, data)
         return last_step
 
     def get_sampler_state(self, training: bool=False) -> dict:


Test patch:
diff --git a/test/integration/test_HMC.py b/test/integration/test_HMC.py
index 567ca78..addf92c 100644
--- a/test/integration/test_HMC.py
+++ b/test/integration/test_HMC.py
@@ -4,47 +4,44 @@
 import jax.numpy as jnp
 from jax.scipy.special import logsumexp
 
-def dual_moon_pe(x):
+def dual_moon_pe(x, data):
     """
     Term 2 and 3 separate the distribution and smear it along the first and second dimension
     """
-    term1 = 0.5 * ((jnp.linalg.norm(x) - 2) / 0.1) ** 2
+    print("compile count")
+    term1 = 0.5 * ((jnp.linalg.norm(x-data) - 2) / 0.1) ** 2
     term2 = -0.5 * ((x[:1] + jnp.array([-3.0, 3.0])) / 0.8) ** 2
     term3 = -0.5 * ((x[1:2] + jnp.array([-3.0, 3.0])) / 0.6) ** 2
     return -(term1 - logsumexp(term2) - logsumexp(term3))
 
+
 n_dim = 5
 n_chains = 15
 n_local_steps = 30
 step_size = 0.1
 n_leapfrog = 10
 
+data = jnp.arange(5)
+
 rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
 initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
 
 HMC = HMC(dual_moon_pe, True, {"step_size": step_size,"n_leapfrog": n_leapfrog, "inverse_metric": jnp.ones(n_dim)})
 
-initial_PE = jax.vmap(HMC.potential)(initial_position)
-
-HMC_kernel = HMC.make_kernel()
+initial_PE = HMC.logpdf_vmap(initial_position, data)
 
-print(HMC_kernel(rng_key_set[0], initial_position[0], initial_PE[0], HMC.params))
-
-HMC_update = HMC.make_update()
-HMC_update = jax.vmap(HMC_update, in_axes = (None, (0, 0, 0, 0, None)), out_axes=(0, 0, 0, 0, None))
+HMC.precompilation(n_chains, n_dim, n_local_steps, data)
 
 initial_position = jnp.repeat(initial_position[:,None], n_local_steps, 1)
 initial_PE = jnp.repeat(initial_PE[:,None], n_local_steps, 1)
 
-state = (rng_key_set[1], initial_position, initial_PE, jnp.zeros((n_chains, n_local_steps,1)), HMC.params)
-
-
-HMC_update(1, state)
+state = (rng_key_set[1], initial_position, initial_PE, jnp.zeros((n_chains, n_local_steps,1)), data, HMC.params)
 
+HMC.update_vmap(1, state)
 HMC_sampler = HMC.make_sampler()
 
-state = HMC_sampler(rng_key_set[1], n_local_steps, initial_position[:, 0])
+state = HMC_sampler(rng_key_set[1], n_local_steps, initial_position[:, 0], data)
 
 
 from flowMC.nfmodel.rqSpline import RQSpline
@@ -69,9 +66,9 @@ def dual_moon_pe(x):
 nf_sampler = Sampler(
     n_dim,
     rng_key_set,
+    jnp.arange(5),
     HMC,
-    dual_moon_pe,
-    model   ,
+    model,
     n_loop_training=n_loop_training,
     n_loop_production=n_loop_production,
     n_local_steps=n_local_steps,
@@ -80,4 +77,4 @@ def dual_moon_pe(x):
     use_global=False,
 )
 
-nf_sampler.sample(initial_position)
\ No newline at end of file
+nf_sampler.sample(initial_position, data)
\ No newline at end of file
diff --git a/test/integration/test_MALA.py b/test/integration/test_MALA.py
index 1584581..985f8da 100644
--- a/test/integration/test_MALA.py
+++ b/test/integration/test_MALA.py
@@ -1,14 +1,16 @@
-from flowMC.sampler.MALA import MALA, mala_sampler_autotune
+from flowMC.sampler.MALA import MALA
 from flowMC.utils.PRNG_keys import initialize_rng_keys
 import jax
 import jax.numpy as jnp
 from jax.scipy.special import logsumexp
+import numpy as np
 
-def dual_moon_pe(x):
+def dual_moon_pe(x, data):
     """
     Term 2 and 3 separate the distribution and smear it along the first and second dimension
     """
-    term1 = 0.5 * ((jnp.linalg.norm(x) - 2) / 0.1) ** 2
+    print("compile count")
+    term1 = 0.5 * ((jnp.linalg.norm(x-data) - 2) / 0.1) ** 2
     term2 = -0.5 * ((x[:1] + jnp.array([-3.0, 3.0])) / 0.8) ** 2
     term3 = -0.5 * ((x[1:2] + jnp.array([-3.0, 3.0])) / 0.6) ** 2
     return -(term1 - logsumexp(term2) - logsumexp(term3))
@@ -19,29 +21,28 @@ def dual_moon_pe(x):
 step_size = 0.01
 n_leapfrog = 10
 
+data = jnp.arange(5)
+
 rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
 initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
 
 MALA_Sampler = MALA(dual_moon_pe, True, {"step_size": step_size})
 
-MALA_Sampler_kernel = MALA_Sampler.make_kernel()
-
+MALA_Sampler.precompilation(n_chains, n_dim, n_local_steps, data)
 
-MALA_Sampler_update = MALA_Sampler.make_update()
-MALA_Sampler_update = jax.vmap(MALA_Sampler_update, in_axes = (None, (0, 0, 0, 0, None)), out_axes=(0, 0, 0, 0, None))
 
 initial_position = jnp.repeat(initial_position[:,None], n_local_steps, 1)
-initial_logp = jnp.repeat(jax.vmap(dual_moon_pe)(initial_position[:,0])[:,None], n_local_steps, 1)[...,None]
+initial_logp = jnp.repeat(jax.vmap(dual_moon_pe, in_axes=(0,None))(initial_position[:,0], data)[:,None], n_local_steps, 1)[...,None]
 
-state = (rng_key_set[1], initial_position, initial_logp, jnp.zeros((n_chains, n_local_steps,1)), MALA_Sampler.params)
+state = (rng_key_set[1], initial_position, initial_logp, jnp.zeros((n_chains, n_local_steps,1)), data, MALA_Sampler.params)
 
 
-MALA_Sampler_update(1, state)
+MALA_Sampler.update_vmap(1, state)
 
 MALA_Sampler_sampler = MALA_Sampler.make_sampler()
 
-state = MALA_Sampler_sampler(rng_key_set[1], n_local_steps, initial_position[:,0])
+state = MALA_Sampler_sampler(rng_key_set[1], n_local_steps, initial_position[:,0], data)
 
 
 from flowMC.nfmodel.rqSpline import RQSpline
@@ -66,18 +67,17 @@ def dual_moon_pe(x):
 nf_sampler = Sampler(
     n_dim,
     rng_key_set,
+    jnp.arange(5),
     MALA_Sampler,
-    dual_moon_pe,
-    model   ,
+    model,
     n_loop_training=n_loop_training,
     n_loop_production=n_loop_production,
     n_local_steps=n_local_steps,
     n_global_steps=n_global_steps,
     n_chains=n_chains,
-    local_autotune=mala_sampler_autotune,
+    # local_autotune=mala_sampler_autotune,
     use_global=False,
 )
 
-nf_sampler.sample(initial_position)
+nf_sampler.sample(initial_position, data)
 
-mala_kernel_vmap = jax.vmap(MALA_Sampler_kernel, in_axes = (0, 0, 0,  None), out_axes=(0, 0, 0))
diff --git a/test/integration/test_RWMCMC.py b/test/integration/test_RWMCMC.py
index 1d5e485..cff32fe 100644
--- a/test/integration/test_RWMCMC.py
+++ b/test/integration/test_RWMCMC.py
@@ -4,11 +4,12 @@
 import jax.numpy as jnp
 from jax.scipy.special import logsumexp
 
-def dual_moon_pe(x):
+def dual_moon_pe(x, data):
     """
     Term 2 and 3 separate the distribution and smear it along the first and second dimension
     """
-    term1 = 0.5 * ((jnp.linalg.norm(x) - 2) / 0.1) ** 2
+    print("compile count")
+    term1 = 0.5 * ((jnp.linalg.norm(x-data) - 2) / 0.1) ** 2
     term2 = -0.5 * ((x[:1] + jnp.array([-3.0, 3.0])) / 0.8) ** 2
     term3 = -0.5 * ((x[1:2] + jnp.array([-3.0, 3.0])) / 0.6) ** 2
     return -(term1 - logsumexp(term2) - logsumexp(term3))
@@ -19,28 +20,26 @@ def dual_moon_pe(x):
 step_size = 0.1
 n_leapfrog = 10
 
+data = jnp.arange(5)
+
 rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
 initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
 
 RWMCMC = GaussianRandomWalk(dual_moon_pe, True, {"step_size": step_size})
 
-RWMCMC_kernel = RWMCMC.make_kernel()
-
-
-RWMCMC_update = RWMCMC.make_update()
-RWMCMC_update = jax.vmap(RWMCMC_update, in_axes = (None, (0, 0, 0, 0, None)), out_axes=(0, 0, 0, 0, None))
+RWMCMC.precompilation(n_chains, n_dim, n_local_steps, data)
 
 initial_position = jnp.repeat(initial_position[:,None], n_local_steps, 1)
-initial_logp = jnp.repeat(jax.vmap(dual_moon_pe)(initial_position[:,0])[:,None], n_local_steps, 1)[...,None]
+initial_logp = jnp.repeat(jax.vmap(dual_moon_pe, in_axes=(0,None))(initial_position[:,0],data)[:,None], n_local_steps, 1)[...,None]
 
-state = (rng_key_set[1], initial_position, initial_logp, jnp.zeros((n_chains, n_local_steps,1)), {"step_size": step_size})
+state = (rng_key_set[1], initial_position, initial_logp, jnp.zeros((n_chains, n_local_steps,1)), data, RWMCMC.params)
 
-RWMCMC_update(1, state)
+RWMCMC.update_vmap(1, state)
 
 RWMCMC_sampler = RWMCMC.make_sampler()
 
-state = RWMCMC_sampler(rng_key_set[1], n_local_steps, initial_position[:,0])
+state = RWMCMC_sampler(rng_key_set[1], n_local_steps, initial_position[:,0], data)
 
 
 from flowMC.nfmodel.rqSpline import RQSpline
@@ -65,9 +64,9 @@ def dual_moon_pe(x):
 nf_sampler = Sampler(
     n_dim,
     rng_key_set,
+    data,
     RWMCMC,
-    dual_moon_pe,
-    model   ,
+    model,
     n_loop_training=n_loop_training,
     n_loop_production=n_loop_production,
     n_local_steps=n_local_steps,
@@ -76,4 +75,4 @@ def dual_moon_pe(x):
     use_global=False,
 )
 
-nf_sampler.sample(initial_position)
\ No newline at end of file
+nf_sampler.sample(initial_position, data)
\ No newline at end of file
diff --git a/test/integration/test_quickstart.py b/test/integration/test_quickstart.py
index 0d401de..a720cb2 100644
--- a/test/integration/test_quickstart.py
+++ b/test/integration/test_quickstart.py
@@ -6,8 +6,10 @@
 from flowMC.utils.PRNG_keys import initialize_rng_keys
 from flowMC.nfmodel.utils import *
 
-def log_posterior(x):
-    return -0.5 * jnp.sum(x ** 2)
+def log_posterior(x, data):
+    return -0.5 * jnp.sum((x-data) ** 2)
+
+data = jnp.arange(5)
 
 n_dim = 5
 n_chains = 10
@@ -20,8 +22,8 @@ def log_posterior(x):
 
 nf_sampler = Sampler(n_dim,
                     rng_key_set,
+                    jnp.arange(n_dim),
                     local_sampler,
-                    log_posterior,
                     model,
                     n_local_steps = 50,
                     n_global_steps = 50,
@@ -30,5 +32,5 @@ def log_posterior(x):
                     batch_size = 1000,
                     n_chains = n_chains)
 
-nf_sampler.sample(initial_position)
+nf_sampler.sample(initial_position, data)
 chains,log_prob,local_accs, global_accs = nf_sampler.get_sampler_state().values()
diff --git a/test/unit/test_kernels.py b/test/unit/test_kernels.py
index 4d95e21..128bb91 100644
--- a/test/unit/test_kernels.py
+++ b/test/unit/test_kernels.py
@@ -1,11 +1,11 @@
 from flowMC.sampler.HMC import HMC
-from flowMC.sampler.MALA import MALA, mala_sampler_autotune
+from flowMC.sampler.MALA import MALA
 from flowMC.sampler.Gaussian_random_walk import GaussianRandomWalk
 from flowMC.utils.PRNG_keys import initialize_rng_keys
 import jax
 import jax.numpy as jnp
 
-def log_posterior(x):
+def log_posterior(x, data=None):
     return -0.5 * jnp.sum(x ** 2)
 
 class TestHMC:
@@ -18,14 +18,14 @@ def test_HMC_deterministic(self):
         rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
         initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
-        initial_PE = jax.vmap(HMC_obj.potential)(initial_position)
+        initial_PE = jax.vmap(HMC_obj.potential)(initial_position, None)
 
         HMC_kernel, leapfrog_kernel, leapfrog_step = HMC_obj.make_kernel(return_aux=True)
 
         # Test whether the HMC kernel is deterministic
 
-        result1 = (HMC_kernel(rng_key_set[0], initial_position[0], initial_PE[0], HMC_obj.params))
-        result2 = (HMC_kernel(rng_key_set[0], initial_position[0], initial_PE[0], HMC_obj.params))
+        result1 = (HMC_kernel(rng_key_set[0], initial_position[0], initial_PE[0], None, HMC_obj.params))
+        result2 = (HMC_kernel(rng_key_set[0], initial_position[0], initial_PE[0], None, HMC_obj.params))
 
         assert jnp.allclose(result1[0],result2[0])
         assert result1[1]==result2[1]
@@ -39,17 +39,17 @@ def test_leapfrog_reversible(self):
 
         rng_key_set = initialize_rng_keys(n_chains, seed=42)
         initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
-        initial_PE = jax.vmap(HMC_obj.potential)(initial_position)
+        initial_PE = jax.vmap(HMC_obj.potential, in_axes=(0,None))(initial_position, None)
 
         HMC_kernel, leapfrog_kernel, leapfrog_step = HMC_obj.make_kernel(return_aux=True)
         key1, key2 = jax.random.split(rng_key_set[0])
 
         initial_momentum = jax.random.normal(key1, shape=initial_position.shape) * jnp.ones(n_dim)**-0.5
-        new_position, new_momentum = leapfrog_step(initial_position, initial_momentum, HMC_obj.params)
-        rev_position, rev_momentum = leapfrog_step(new_position, -new_momentum, HMC_obj.params)
+        new_position, new_momentum = leapfrog_step(initial_position, initial_momentum, None, HMC_obj.params)
+        rev_position, rev_momentum = leapfrog_step(new_position, -new_momentum, None, HMC_obj.params)
 
         assert jnp.allclose(rev_position, initial_position)
-        assert jnp.allclose(initial_PE, HMC_obj.potential(rev_position))
+        assert jnp.allclose(initial_PE, HMC_obj.potential(rev_position, None))
 
     def test_HMC_acceptance_rate(self):
         # Test acceptance rate goes to one when step size is small
@@ -62,9 +62,9 @@ def test_HMC_acceptance_rate(self):
         rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
         initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
-        initial_PE = jax.vmap(HMC_obj.potential)(initial_position)
+        initial_PE = jax.vmap(HMC_obj.potential)(initial_position, None)
 
-        result = jax.vmap(HMC_kernel, in_axes = (0, 0, 0, None), out_axes=(0, 0, 0))(rng_key_set[1], initial_position, initial_PE, HMC_obj.params)
+        result = jax.vmap(HMC_kernel, in_axes = (0, 0, 0, None, None), out_axes=(0, 0, 0))(rng_key_set[1], initial_position, initial_PE, None, HMC_obj.params)
 
         assert result[2].all()
 
@@ -76,10 +76,10 @@ def test_HMC_close_gaussian(self):
         rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
         initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
-
+        HMC_obj.precompilation(n_chains, n_dim, 10000, None)
         HMC_sampler = HMC_obj.make_sampler()
 
-        result = HMC_sampler(rng_key_set[1], 10000, initial_position)
+        result = HMC_sampler(rng_key_set[1], 10000, initial_position, None)
 
         assert jnp.isclose(jnp.mean(result[1]),0,atol=1e-2)
         assert jnp.isclose(jnp.var(result[1]),1,atol=1e-2)
@@ -95,7 +95,7 @@ def test_MALA_deterministic(self):
         rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
         initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
-        initial_logp = log_posterior(initial_position)
+        initial_logp = log_posterior(initial_position, None)
 
         result1 = (MALA_kernel(rng_key_set[0], initial_position[0], initial_logp, MALA_obj.params))
         result2 = (MALA_kernel(rng_key_set[0], initial_position[0], initial_logp, MALA_obj.params))
@@ -115,9 +115,9 @@ def test_MALA_acceptance_rate(self):
         rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
         initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
-        initial_logp = jax.vmap(log_posterior)(initial_position)
+        initial_logp = jax.vmap(log_posterior)(initial_position, None)
 
-        result = jax.vmap(MALA_kernel, in_axes = (0, 0, 0, None), out_axes=(0, 0, 0))(rng_key_set[1], initial_position, initial_logp, MALA_obj.params)
+        result = jax.vmap(MALA_kernel, in_axes = (0, 0, 0, None, None), out_axes=(0, 0, 0))(rng_key_set[1], initial_position, initial_logp, None, MALA_obj.params)
 
         assert result[2].all()
 
@@ -129,10 +129,11 @@ def test_MALA_close_gaussian(self):
         rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
         initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
+        MALA_obj.precompilation(n_chains, n_dim, 30000, None)
 
         MALA_sampler = MALA_obj.make_sampler()
 
-        result = MALA_sampler(rng_key_set[1], 30000, initial_position)
+        result = MALA_sampler(rng_key_set[1], 30000, initial_position, None)
 
         assert jnp.isclose(jnp.mean(result[1]),0,atol=1e-2)
         assert jnp.isclose(jnp.var(result[1]),1,atol=1e-2)
@@ -171,7 +172,7 @@ def test_Gaussian_random_walk_acceptance_rate(self):
         initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
         initial_logp = jax.vmap(log_posterior)(initial_position)
 
-        result = jax.vmap(GRW_kernel, in_axes = (0, 0, 0, None), out_axes=(0, 0, 0))(rng_key_set[1], initial_position, initial_logp, GRW_obj.params)
+        result = jax.vmap(GRW_kernel, in_axes = (0, 0, 0, None, None), out_axes=(0, 0, 0))(rng_key_set[1], initial_position, initial_logp, None, GRW_obj.params)
 
         assert result[2].all()
 
@@ -183,10 +184,10 @@ def test_Gaussian_random_walk_close_gaussian(self):
         rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
         initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
-
+        GRW_obj.precompilation(n_chains, n_dim, 30000, None)
         GRW_sampler = GRW_obj.make_sampler()
 
-        result = GRW_sampler(rng_key_set[1], 30000, initial_position)
+        result = GRW_sampler(rng_key_set[1], 30000, initial_position, None)
 
         assert jnp.isclose(jnp.mean(result[1]),0,atol=1e-2)
         assert jnp.isclose(jnp.var(result[1]),1,atol=1e-2)
\ No newline at end of file