Repository: bb515/diffusionjax (language: python)

Issue:
Refactor/pylintrc



Discussion / hints:

Fix patch:
diff --git a/.editorconfig b/.editorconfig
new file mode 100644
index 0000000..f5607c0
--- /dev/null
+++ b/.editorconfig
@@ -0,0 +1,4 @@
+# 2 space indentation
+[*.py]
+indent_style = space
+indent_size = 2
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 2fe461d..e23ecda 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -19,7 +19,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip install --upgrade pip setuptools numpy
-          pip install --no-cache-dir -U -r requirements.txt | cat
+          pip install -e '.[linting,testing]'
           pip install --upgrade numpy
       - name: Test
         env:
diff --git a/README.md b/README.md
index 667954e..92f6858 100644
--- a/README.md
+++ b/README.md
@@ -2,7 +2,6 @@ diffusionjax
 ============
 [![CI](https://github.com/bb515/diffusionjax/actions/workflows/CI.yml/badge.svg)](https://github.com/bb515/diffusionjax/actions/workflows/CI.yml)
 [![Coverage Status](https://coveralls.io/repos/github/bb515/diffusionjax/badge.svg?branch=master)](https://coveralls.io/github/bb515/diffusionjax?branch=master)
-[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 
 diffusionjax is a simple, accessible introduction to diffusion models, also known as score-based generative models (SGMs). It is implemented in Python via the autodiff framework, [JAX](https://github.com/google/jax). In particular, diffusionjax uses the [Flax](https://github.com/google/flax) library for the neural network approximator of the score.
 
@@ -13,7 +12,6 @@ The development of diffusionjax has been supported by The Alan Turing Institute
 Thank you to [nPlan](https://www.nplan.io/), who are supporting this project.
 
 Contents:
-
 - [Installation](#installation)
 - [Examples](#examples)
     - [Introduction to diffusion models](#introduction-to-diffusion-models)
@@ -25,6 +23,7 @@ Contents:
 The package requires Python 3.8+. First, it is recommended to [create a new python virtual environment](https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#creating-an-environment-with-commands). 
 diffusionjax depends on JAX. Because the JAX installation is different depending on your CUDA version, diffusionjax does not list JAX as a dependency in `setup.py`.
 First, [follow these instructions](https://github.com/google/jax#installation) to install JAX with the relevant accelerator support.
+To run the examples, you may optionally need to install [optax](https://optax.readthedocs.io/en/latest/), [orbax-checkpoint](https://orbax.readthedocs.io/en/latest/), [torch[cpu]](https://pytorch.org/get-started/locally/) and [mlkernels](https://github.com/wesselb/mlkernels#installation), which the package depends on only through the examples given.
 Then, `pip install diffusionjax` or for developers,
 - Clone the repository `git clone git@github.com:bb515/diffusionjax.git`
 - Install using pip `pip install -e .` from the root directory of the repository (see the `setup.py` for the requirements that this command installs).
diff --git a/diffusionjax/inverse_problems.py b/diffusionjax/inverse_problems.py
new file mode 100644
index 0000000..588a0b7
--- /dev/null
+++ b/diffusionjax/inverse_problems.py
@@ -0,0 +1,132 @@
+"""Utility functions related to Bayesian inversion."""
+import jax.numpy as jnp
+from jax.lax import scan
+from jax import vmap
+import jax.random as random
+from diffusionjax.utils import batch_mul
+
+
+def merge_data_with_mask(x_space, data, mask, coeff=1.):
+  return data * mask * coeff + x_space * (1. - mask * coeff)
+
+
+def get_projection_sampler(solver, inverse_scaler=None, denoise=True, stack_samples=False):
+  """Create an image inpainting function that uses sampler, that returns only the final sample.
+
+  Args:
+    inverse_scaler: The inverse data normalizer.
+    denoise: Boolean variable that if `True` applies one-step denoising to final samples.
+    stack_samples: Boolean variable that if `True` returns all samples on path(s).
+  Returns:
+    A projection sampler function.
+  """
+  if inverse_scaler is None:
+    def inverse_scaler(x):
+      return x
+  vmap_inverse_scaler = vmap(inverse_scaler)
+
+  def update(rng, data, mask, x, vec_t, coeff):
+    data_mean, std = solver.sde.marginal_prob(data, vec_t)
+    z = random.normal(rng, x.shape)
+    z_data = data_mean + batch_mul(std, z)
+    x = merge_data_with_mask(x, z_data, mask, coeff)
+    rng, step_rng = random.split(rng)
+    x, x_mean = solver.update(step_rng, x, vec_t)
+    return x, x_mean
+
+  ts = solver.ts
+  num_function_evaluations = jnp.size(ts)
+
+  def projection_sampler(rng, data, mask, coeff):
+    """Sampler for image inpainting.
+
+    Args:
+      rng: A JAX random state.
+      data: A JAX array that represents a mini-batch of images to inpaint.
+      mask: A {0, 1} array with the same shape of `data`. Value `1` marks known pixels,
+        and value `0` marks pixels that require inpainting.
+
+    Returns:
+      Inpainted (complete) images.
+    """
+    # Initial sample
+    rng, step_rng = random.split(rng)
+    shape = data.shape
+    x = random.normal(step_rng, shape)
+
+    def f(carry, t):
+      rng, x, x_mean = carry
+      vec_t = jnp.full(shape[0], t)
+      rng, step_rng = random.split(rng)
+      x, x_mean = update(step_rng, data, mask, x, vec_t, coeff)
+      return ((rng, x, x_mean), x) if stack_samples else ((rng, x, x_mean), ())
+
+    if not stack_samples:
+      (_, x, _), _ = scan(f, (rng, x, x), ts, reverse=True)
+      return inverse_scaler(x), num_function_evaluations
+    else:
+      (_, _, _), xs = scan(f, (rng, x, x), ts, reverse=True)
+      return vmap_inverse_scaler(xs), num_function_evaluations
+
+  # return jax.pmap(projection_sampler, in_axes=(0, None, None, None), axis_name='batch')
+  return projection_sampler
+
+
+def get_inpainter(solver, inverse_scaler=None, stack_samples=False):
+  """Create an image inpainting function that uses sampler, that returns only the final sample.
+
+  Args:
+    inverse_scaler: The inverse data normalizer.
+    denoise: Boolean variable that if `True` applies one-step denoising to final samples.
+  Returns:
+    An inpainting function.
+  """
+  if inverse_scaler is None:
+    def inverse_scaler(x):
+      return x
+  vmap_inverse_scaler = vmap(inverse_scaler)
+
+  def update(rng, data, mask, x, vec_t):
+    x, x_mean = solver.update(rng, x, vec_t)
+    masked_data_mean, std = solver.sde.marginal_prob(data, vec_t)
+    rng, step_rng = random.split(rng)
+    masked_data = masked_data_mean + batch_mul(random.normal(step_rng, x.shape), std)
+    x = x * (1. - mask) + masked_data * mask
+    x_mean = x * (1. - mask) + masked_data_mean * mask
+    return x, x_mean
+
+  ts = solver.ts
+  num_function_evaluations = jnp.size(ts)
+
+  def inpainter(rng, data, mask):
+    """Sampler for image inpainting.
+
+    Args:
+      rng: A JAX random state.
+      data: A JAX array that represents a mini-batch of images to inpaint.
+      mask: A {0, 1} array with the same shape of `data`. Value `1` marks known pixels,
+        and value `0` marks pixels that require inpainting.
+
+    Returns:
+      Inpainted (complete) images.
+    """
+    # Initial sample
+    rng, step_rng = random.split(rng)
+    shape = data.shape
+    x = data * mask + random.normal(step_rng, shape) * (1. - mask)
+
+    def f(carry, t):
+      rng, x, x_mean = carry
+      vec_t = jnp.full(shape[0], t)
+      rng, step_rng = random.split(rng)
+      x, x_mean = update(rng, data, mask, x, vec_t)
+      return ((rng, x, x_mean), x) if stack_samples else ((rng, x, x_mean), ())
+
+    if not stack_samples:
+      (_, x, _), _ = scan(f, (rng, x, x), ts, reverse=True)
+      return inverse_scaler(x), num_function_evaluations
+    else:
+      (_, _, _), xs = scan(f, (rng, x, x), ts, reverse=True)
+      return vmap_inverse_scaler(xs), num_function_evaluations
+  # return jax.pmap(inpainter, in_axes=(0, None, None), axis_name='batch')
+  return inpainter
\ No newline at end of file
diff --git a/diffusionjax/models.py b/diffusionjax/models.py
index aac17dc..5e972da 100644
--- a/diffusionjax/models.py
+++ b/diffusionjax/models.py
@@ -1,54 +1,47 @@
+"""Functions are designed for a mini-batch of inputs."""
 import flax.linen as nn
 import numpy as np
 import jax.numpy as jnp
 
 
 class MLP(nn.Module):
-    """
-    A simple model with multiple fully connected layers and some fourier features
-    for the time variable. Functions are designed for a mini-batch of inputs.
-    """
-    @nn.compact
-    def __call__(self, x, t):
-        x_shape = x.shape
-        in_size = np.prod(x_shape[1:])
-        n_hidden = 256
-        t = t.reshape((t.shape[0], -1))
-        x = x.reshape((x.shape[0], -1))  # flatten
-        t = jnp.concatenate([t - 0.5, jnp.cos(2*jnp.pi*t)], axis=-1)
-        x = jnp.concatenate([x, t], axis=-1)
-        x = nn.Dense(n_hidden)(x)
-        x = nn.relu(x)
-        x = nn.Dense(n_hidden)(x)
-        x = nn.relu(x)
-        x = nn.Dense(n_hidden)(x)
-        x = nn.relu(x)
-        x = nn.Dense(in_size)(x)
-        return x.reshape(x_shape)
+  @nn.compact
+  def __call__(self, x, t):
+    x_shape = x.shape
+    in_size = np.prod(x_shape[1:])
+    n_hidden = 256
+    t = t.reshape((t.shape[0], -1))
+    x = x.reshape((x.shape[0], -1))  # flatten
+    t = jnp.concatenate([t - 0.5, jnp.cos(2*jnp.pi*t)], axis=-1)
+    x = jnp.concatenate([x, t], axis=-1)
+    x = nn.Dense(n_hidden)(x)
+    x = nn.relu(x)
+    x = nn.Dense(n_hidden)(x)
+    x = nn.relu(x)
+    x = nn.Dense(n_hidden)(x)
+    x = nn.relu(x)
+    x = nn.Dense(in_size)(x)
+    return x.reshape(x_shape)
 
 
 class CNN(nn.Module):
-    """
-    A simple model with a single convolutional layer with time as another channel.
-    Functions are designed for a mini-batch of inputs.
-    """
-    @nn.compact
-    def __call__(self, x, t):
-        x_shape = x.shape
-        ndim = x.ndim
+  @nn.compact
+  def __call__(self, x, t):
+    x_shape = x.shape
+    ndim = x.ndim
 
-        n_hidden = x_shape[1]
-        n_time_channels = 1
+    n_hidden = x_shape[1]
+    n_time_channels = 1
 
-        t = t.reshape((t.shape[0], -1))
-        t = jnp.concatenate([t - 0.5, jnp.cos(2*jnp.pi*t)], axis=-1)
-        t = nn.Dense(n_hidden**2 * n_time_channels)(t)
-        t = nn.relu(t)
-        t = nn.Dense(n_hidden**2 * n_time_channels)(t)
-        t = nn.relu(t)
-        t = t.reshape(t.shape[0], n_hidden, n_hidden, n_time_channels)
-        # Add time as another channel
-        x = jnp.concatenate((x, t), axis=-1)
-        # Global convolution
-        x = nn.Conv(x_shape[-1], kernel_size=(9,) * (ndim - 2))(x)
-        return x
+    t = t.reshape((t.shape[0], -1))
+    t = jnp.concatenate([t - 0.5, jnp.cos(2*jnp.pi*t)], axis=-1)
+    t = nn.Dense(n_hidden**2 * n_time_channels)(t)
+    t = nn.relu(t)
+    t = nn.Dense(n_hidden**2 * n_time_channels)(t)
+    t = nn.relu(t)
+    t = t.reshape(t.shape[0], n_hidden, n_hidden, n_time_channels)
+    # Add time as another channel
+    x = jnp.concatenate((x, t), axis=-1)
+    # A single convolution layer
+    x = nn.Conv(x_shape[-1], kernel_size=(9,) * (ndim - 2))(x)
+    return x
diff --git a/diffusionjax/plot.py b/diffusionjax/plot.py
index f9cc547..59da616 100644
--- a/diffusionjax/plot.py
+++ b/diffusionjax/plot.py
@@ -1,177 +1,171 @@
-"""Helper functions for plots in the example."""
+"""Plotting code for the examples."""
 import jax.numpy as jnp
 import matplotlib.pyplot as plt
 from jax import jit, vmap
 from functools import partial
-import seaborn as sns
-sns.set_style("darkgrid")
-cm = sns.color_palette("mako_r", as_cmap=True)
 import matplotlib.animation as animation
 
+
 BG_ALPHA = 1.0
 MG_ALPHA = 0.2
 FG_ALPHA = 0.4
 
 
-def plot_heatmap(samples, area_min=-3, area_max=3, lengthscale=350, fname="plot_heatmap"):
-    """Plots a heatmap of all samples in the area [area_min, area_max] x [area_min, area_max].
-    Args:
-        samples: locations of all particles in R^2, array (J, 2)
-        area_min: lowest x and y coordinate
-        area_max: highest x and y coordinate
-    """
-    def small_kernel(z, area_min, area_max):
-        a = jnp.linspace(area_min, area_max, 512)
-        x, y = jnp.meshgrid(a, a)
-        dist = (x - z[0])**2 + (y - z[1])**2
-        hm = jnp.exp(-lengthscale * dist)
-        return hm
-
-    # We try to jit most of the code, but use the helper functions
-    # since we cannot jit all of it because of the plt functions
-    @jit
-    def produce_heatmap(samples, area_min, area_max):
-        return jnp.sum(vmap(small_kernel, in_axes=(0, None, None))(samples, area_min, area_max), axis=0)
-
-    hm = produce_heatmap(samples, area_min, area_max)
-    extent = [area_min, area_max, area_max, area_min]
-    plt.imshow(hm, cmap=cm, interpolation='nearest', extent=extent)
-    ax = plt.gca()
-    ax.invert_yaxis()
-    plt.savefig(fname)
-    plt.close()
+def plot_heatmap(samples, area_bounds, lengthscale=350.0, fname="plot_heatmap") -> None:
+  """Plots a heatmap of all samples in the area area_bounds x area_bounds.
+  Args:
+    samples: locations of particles shape (num_particles, 2)
+  """
+  def small_kernel(z, area_bounds):
+    a = jnp.linspace(area_bounds[0], area_bounds[1], 512)
+    x, y = jnp.meshgrid(a, a)
+    dist = (x - z[0])**2 + (y - z[1])**2
+    hm = jnp.exp(-lengthscale * dist)
+    return hm
+
+  # jit most of the code, but use the helper functions since cannot jit all of it because of plt
+  @jit
+  def produce_heatmap(samples, area_bounds):
+    return jnp.sum(vmap(small_kernel, in_axes=(0, None))(samples, area_bounds), axis=0)
+
+  hm = produce_heatmap(samples, area_bounds)
+  extent = area_bounds + area_bounds
+  plt.imshow(hm, interpolation='nearest', extent=extent)
+  ax = plt.gca()
+  ax.invert_yaxis()
+  plt.savefig(fname)
+  plt.close()
 
 
 def plot_samples(samples, index, fname="samples.png", lims=None):
-    fig, ax = plt.subplots(1, 1)
-    fig.patch.set_facecolor('white')
-    fig.patch.set_alpha(BG_ALPHA)
-    ax.scatter(
-        samples[:, index[0]], samples[:, index[1]],
-        color='red', label=r"$x$")
-    ax.legend()
-    ax.set_xlabel(r"$x_{}$".format(index[0]))
-    ax.set_ylabel(r"$x_{}$".format(index[1]))
-    if lims is not None:
-        ax.set_xlim(lims[0])
-        ax.set_ylim(lims[1])
-    plt.gca().set_aspect('equal', adjustable='box')
-    plt.draw()
-    fig.savefig(
-        fname,
-        facecolor=fig.get_facecolor(), edgecolor='none')
-    plt.close()
+  fig, ax = plt.subplots(1, 1)
+  fig.patch.set_facecolor('white')
+  fig.patch.set_alpha(BG_ALPHA)
+  ax.scatter(
+    samples[:, index[0]], samples[:, index[1]],
+    color='red', label=r"$x$")
+  ax.legend()
+  ax.set_xlabel(r"$x_{}$".format(index[0]))
+  ax.set_ylabel(r"$x_{}$".format(index[1]))
+  if lims is not None:
+    ax.set_xlim(lims[0])
+    ax.set_ylim(lims[1])
+  plt.gca().set_aspect('equal', adjustable='box')
+  plt.draw()
+  fig.savefig(
+    fname,
+    facecolor=fig.get_facecolor(), edgecolor='none')
+  plt.close()
 
 
 def plot_animation(fig, ax, animate, frames, fname, fps=20, bitrate=800, dpi=300):
-
-    ani = animation.FuncAnimation(
-        fig, animate, frames=frames, interval=1, fargs=(ax,))
-    # Set up formatting for the movie files
-    Writer = animation.writers['ffmpeg']
-    writer = Writer(fps=fps, metadata=dict(artist='Me'), bitrate=bitrate)
-    # Note that mp4 does not work on pdf
-    ani.save('{}.mp4'.format(fname), writer=writer, dpi=dpi)
+  ani = animation.FuncAnimation(
+    fig, animate, frames=frames, interval=1, fargs=(ax,))
+  # Set up formatting for the movie files
+  Writer = animation.writers['ffmpeg']
+  writer = Writer(fps=fps, metadata=dict(artist='Me'), bitrate=bitrate)
+  # Note that mp4 does not work on pdf
+  ani.save('{}.mp4'.format(fname), writer=writer, dpi=dpi)
 
 
 def plot_score(score, scaler, t, area_min=-1, area_max=1, fname="plot_score"):
-    fig, ax = plt.subplots(1, 1)
-    # this helper function is here so that we can jit it.
-    # We can not jit the whole function since plt.quiver cannot
-    # be jitted
-    @partial(jit, static_argnums=[0,])
-    def helper(score, t, area_min, area_max):
-        x = jnp.linspace(area_min, area_max, 16)
-        x, y = jnp.meshgrid(x, x)
-        grid = jnp.stack([x.flatten(), y.flatten()], axis=1)
-        t = jnp.ones((grid.shape[0],)) * t
-        scores = score(scaler(grid), t)
-        return grid, scores
-    grid, scores = helper(score, t, area_min, area_max)
-    ax.quiver(grid[:, 0], grid[:, 1], scores[:, 0], scores[:, 1])
-    ax.set_xlabel(r"$x_0$")
-    ax.set_ylabel(r"$x_1$")
-    plt.gca().set_aspect('equal', adjustable='box')
-    fig.savefig(fname)
-    plt.close()
+  fig, ax = plt.subplots(1, 1)
+  # this helper function is here so that we can jit it.
+  # We can not jit the whole function since plt.quiver cannot
+  # be jitted
+  @partial(jit, static_argnums=[0,])
+  def helper(score, t, area_min, area_max):
+    x = jnp.linspace(area_min, area_max, 16)
+    x, y = jnp.meshgrid(x, x)
+    grid = jnp.stack([x.flatten(), y.flatten()], axis=1)
+    t = jnp.ones((grid.shape[0],)) * t
+    scores = score(scaler(grid), t)
+    return grid, scores
+  grid, scores = helper(score, t, area_min, area_max)
+  ax.quiver(grid[:, 0], grid[:, 1], scores[:, 0], scores[:, 1])
+  ax.set_xlabel(r"$x_0$")
+  ax.set_ylabel(r"$x_1$")
+  plt.gca().set_aspect('equal', adjustable='box')
+  fig.savefig(fname)
+  plt.close()
 
 
 def plot_score_ax(ax, score, scaler, t, area_min=-1, area_max=1):
-    # This helper function is here so that we can jit it.
-    # We can not jit the whole function since plt.quiver cannot be jitted
-    @partial(jit, static_argnums=[0,])
-    def helper(score, t, area_min, area_max):
-        x = jnp.linspace(area_min, area_max, 16)
-        x, y = jnp.meshgrid(x, x)
-        grid = jnp.stack([x.flatten(), y.flatten()], axis=1)
-        t = jnp.ones((grid.shape[0],)) * t
-        scores = score(scaler(grid), t)
-        return grid, scores
-    grid, scores = helper(score, t, area_min, area_max)
-    ax.quiver(grid[:, 0], grid[:, 1], scores[:, 0], scores[:, 3])
-    ax.set_xlabel(r"$x_0$")
-    ax.set_ylabel(r"$x_1$")
+  # This helper function is here so that we can jit it.
+  # We can not jit the whole function since plt.quiver cannot be jitted
+  @partial(jit, static_argnums=[0,])
+  def helper(score, t, area_min, area_max):
+    x = jnp.linspace(area_min, area_max, 16)
+    x, y = jnp.meshgrid(x, x)
+    grid = jnp.stack([x.flatten(), y.flatten()], axis=1)
+    t = jnp.ones((grid.shape[0],)) * t
+    scores = score(scaler(grid), t)
+    return grid, scores
+  grid, scores = helper(score, t, area_min, area_max)
+  ax.quiver(grid[:, 0], grid[:, 1], scores[:, 0], scores[:, 3])
+  ax.set_xlabel(r"$x_0$")
+  ax.set_ylabel(r"$x_1$")
 
 
 def plot_heatmap_ax(ax, samples, area_min=-3, area_max=3, lengthscale=350):
-    """Plots a heatmap of all samples in the area [area_min, area_max] x [area_min, area_max].
-    Args:
-        samples: locations of all particles in R^2, array (J, 2)
-        area_min: lowest x and y coordinate
-        area_max: highest x and y coordinate
-    """
-    def small_kernel(z, area_min, area_max):
-        a = jnp.linspace(area_min, area_max, 512)
-        x, y = jnp.meshgrid(a, a)
-        dist = (x - z[0])**2 + (y - z[1])**2
-        hm = jnp.exp(-lengthscale * dist)
-        return hm
-
-    # We try to jit most of the code, but use the helper functions
-    # since we cannot jit all of it because of the plt functions
-    @jit
-    def produce_heatmap(samples, area_min, area_max):
-        return jnp.sum(vmap(small_kernel, in_axes=(0, None, None))(samples, area_min, area_max), axis=0)
-
-    hm = produce_heatmap(samples, area_min, area_max)
-    extent = [area_min, area_max, area_max, area_min]
-    ax.imshow(hm, cmap=cm, interpolation='nearest', extent=extent)
-    ax = plt.gca()
-    ax.invert_yaxis()
-    ax.set_xlabel(r"$x_0$")
-    ax.set_ylabel(r"$x_1$")
+  """Plots a heatmap of all samples in the area [area_min, area_max] x [area_min, area_max].
+  Args:
+    samples: locations of all particles in R^2, array (J, 2)
+    area_min: lowest x and y coordinate
+    area_max: highest x and y coordinate
+  """
+  def small_kernel(z, area_min, area_max):
+    a = jnp.linspace(area_min, area_max, 512)
+    x, y = jnp.meshgrid(a, a)
+    dist = (x - z[0])**2 + (y - z[1])**2
+    hm = jnp.exp(-lengthscale * dist)
+    return hm
+
+  # We try to jit most of the code, but use the helper functions
+  # since we cannot jit all of it because of the plt functions
+  @jit
+  def produce_heatmap(samples, area_min, area_max):
+    return jnp.sum(vmap(small_kernel, in_axes=(0, None, None))(samples, area_min, area_max), axis=0)
+
+  hm = produce_heatmap(samples, area_min, area_max)
+  extent = [area_min, area_max, area_max, area_min]
+  ax.imshow(hm, interpolation='nearest', extent=extent)
+  ax = plt.gca()
+  ax.invert_yaxis()
+  ax.set_xlabel(r"$x_0$")
+  ax.set_ylabel(r"$x_1$")
 
 
 def plot_temperature_schedule(sde, solver):
-    """Plots the temperature schedule of the SDE marginals.
+  """Plots the temperature schedule of the SDE marginals.
 
-    Args:
-        sde: a valid SDE class.
-    """
-    m2 = sde.mean_coeff(solver.ts)**2
-    v = sde.variance(solver.ts)
-    plt.plot(solver.ts, m2, label="m2")
-    plt.plot(solver.ts, v, label="v")
-    plt.legend()
-    plt.savefig("plot_temperature_schedule.png")
-    plt.close()
+  Args:
+    sde: a valid SDE class.
+  """
+  m2 = sde.mean_coeff(solver.ts)**2
+  v = sde.variance(solver.ts)
+  plt.plot(solver.ts, m2, label="m2")
+  plt.plot(solver.ts, v, label="v")
+  plt.legend()
+  plt.savefig("plot_temperature_schedule.png")
+  plt.close()
 
 
 def plot_scatter(samples, fname="samples"):
-    fig, ax = plt.subplots(1, 1)
-    fig.patch.set_facecolor('white')
-    fig.patch.set_alpha(1.0)
-    ax.scatter(
-        samples[:, 0], samples[:, 1],
-        alpha=0.1, label=r"$x$")
-    ax.legend()
-    ax.set_xlabel(r"$x_{}$".format(0))
-    ax.set_ylabel(r"$x_{}$".format(1))
-    ax.set_xlim(-3, 3)
-    ax.set_ylim(-3, 3)
-    plt.gca().set_aspect('equal', adjustable='box')
-    plt.draw()
-    fig.savefig(
-        fname,
-        facecolor=fig.get_facecolor(), edgecolor='none')
-    plt.close()
+  fig, ax = plt.subplots(1, 1)
+  fig.patch.set_facecolor('white')
+  fig.patch.set_alpha(1.0)
+  ax.scatter(
+    samples[:, 0], samples[:, 1],
+    alpha=0.1, label=r"$x$")
+  ax.legend()
+  ax.set_xlabel(r"$x_{}$".format(0))
+  ax.set_ylabel(r"$x_{}$".format(1))
+  ax.set_xlim(-3, 3)
+  ax.set_ylim(-3, 3)
+  plt.gca().set_aspect('equal', adjustable='box')
+  plt.draw()
+  fig.savefig(
+    fname,
+    facecolor=fig.get_facecolor(), edgecolor='none')
+  plt.close()
diff --git a/diffusionjax/run_lib.py b/diffusionjax/run_lib.py
index 1a5bd00..73e9ba5 100644
--- a/diffusionjax/run_lib.py
+++ b/diffusionjax/run_lib.py
@@ -1,500 +1,519 @@
-"""
-Training and evaluation for score-based generative models.
-"""
+"""Training and evaluation for score-based generative models."""
 import jax
-from jax import jit
+from jax import jit, value_and_grad
 import jax.random as random
 import jax.numpy as jnp
-
-from diffusionjax.utils import batch_mul, get_loss, get_score, get_sampler
+from diffusionjax.utils import get_loss, get_score, get_sampler
 from diffusionjax.models import MLP, CNN
 import diffusionjax.sde as sde_lib
 from diffusionjax.solvers import EulerMaruyama, Annealed, DDIMVP, DDIMVE, SMLD, DDPM
-from diffusionjax.utils import get_step_fn
-
-from torch.utils.data import DataLoader
 import numpy as np
-import optax
 from functools import partial
 import flax
+import flax.training.orbax_utils as orbax_utils
 import flax.jax_utils as flax_utils
-from flax.training import checkpoints
 from absl import flags
 from tqdm import tqdm, trange
-import gc
-import io
 import os
 import time
 from typing import Any
-import datetime
 import logging
-from collections import defaultdict
 import wandb
-import orbax
+
+
+# This file requires optax, https://optax.readthedocs.io/en/latest/
+import optax
+# This file requires orbax, https://orbax.readthedocs.io/en/latest/
+import orbax.checkpoint
+# This file requires torch[cpu], https://pytorch.org/get-started/locally/
+from torch.utils.data import DataLoader
 
 FLAGS = flags.FLAGS
 logger = logging.getLogger(__name__)
 
 
+def get_step_fn(loss, optimizer, train, pmap):
+  """Create a one-step training/evaluation function.
+
+  Args:
+    loss: A loss function.
+    optimizer: An optimization function.
+    train: `True` for training and `False` for evaluation.
+    pmap: `True` for pmap across jax devices, `False` for single device.
+
+  Returns:
+    A one-step function for training or evaluation.
+  """
+  @jit
+  def step_fn(carry, batch):
+    (rng, params, opt_state) = carry
+    rng, step_rng = random.split(rng)
+    grad_fn = value_and_grad(loss)
+    if train:
+      loss_val, grads = grad_fn(params, step_rng, batch)
+      if pmap:
+        loss_val = jax.lax.pmean(loss_val, axis_name='batch')
+        grads = jax.lax.pmean(grads, axis_name='batch')
+      updates, opt_state = optimizer.update(grads, opt_state)
+      params = optax.apply_updates(params, updates)
+    else:
+      loss_val = loss(params, step_rng, batch)
+      if pmap: loss_val = jax.lax.pmean(loss_val, axis_name='batch')
+    return (rng, params, opt_state), loss_val
+  return step_fn
+
+
 # The dataclass that stores all training states
 @flax.struct.dataclass
 class State:
-    step: int
-    opt_state: Any
-    params: Any
-    rng: Any
-    lr: float
+  step: int
+  opt_state: Any
+  params: Any
+  rng: Any
+  lr: float
 
 
 def get_sde(config):
-    # Setup SDE
-    if config.training.sde.lower()=='vpsde':
-        sde = sde_lib.VP(beta_min=config.model.beta_min, beta_max=config.model.beta_max)
-    elif config.training.sde.lower()=='vesde':
-        sde = sde_lib.VE(sigma_min=config.model.sigma_min, sigma_max=config.model.sigma_max)
-    else:
-        raise NotImplementedError(f"SDE {config.training.SDE} unknown.")
-    return sde
+  # Setup SDE
+  if config.training.sde.lower()=='vpsde':
+    sde = sde_lib.VP(beta_min=config.model.beta_min, beta_max=config.model.beta_max)
+  elif config.training.sde.lower()=='vesde':
+    sde = sde_lib.VE(sigma_min=config.model.sigma_min, sigma_max=config.model.sigma_max)
+  else:
+    raise NotImplementedError(f"SDE {config.training.SDE} unknown.")
+  return sde
 
 
 def get_optimizer(config):
-    """Returns an optax optimizer object based on `config`."""
-    if config.optim.warmup:
-        schedule = optax.warmup_cosine_decay_schedule(
-            init_value=0.0,
-            peak_value=1.0,
-            warmup_steps=config.optim.warmup,
-            decay_steps=config.optim.warmup + 1,
-            end_value=1.0,
-            )
-    else:
-        schedule = config.optim.lr
-    if config.optim.optimizer=='Adam':
-        if config.optim.weight_decay:
-            optimizer = optax.adamw(
-                learning_rate=schedule, b1=config.optim.beta1, eps=config.optim.eps)
-        else:
-            optimizer = optax.adam(
-                learning_rate=schedule, b1=config.optim.beta1, eps=config.optim.eps)
+  """Returns an optax optimizer object based on `config`."""
+  if config.optim.warmup:
+    schedule = optax.warmup_cosine_decay_schedule(
+      init_value=0.0,
+      peak_value=1.0,
+      warmup_steps=config.optim.warmup,
+      decay_steps=config.optim.warmup + 1,
+      end_value=1.0,
+      )
+  else:
+    schedule = config.optim.lr
+  if config.optim.optimizer=='Adam':
+    if config.optim.weight_decay:
+      optimizer = optax.adamw(
+        learning_rate=schedule, b1=config.optim.beta1, eps=config.optim.eps)
     else:
-        raise NotImplementedError(
-            'Optimiser {} not supported yet!'.format(config.optim.optimizer)
-        )
-    if config.optim.grad_clip:
-        optimizer = optax.chain(
-            optax.clip(config.optim.grad_clip),
-            optimizer
-        )
-    return optimizer
+      optimizer = optax.adam(
+        learning_rate=schedule, b1=config.optim.beta1, eps=config.optim.eps)
+  else:
+      raise NotImplementedError(
+        'Optimiser {} not supported yet!'.format(config.optim.optimizer)
+      )
+  if config.optim.grad_clip:
+    optimizer = optax.chain(
+      optax.clip(config.optim.grad_clip),
+      optimizer
+    )
+  return optimizer
 
 
 def get_model(config):
-    if config.model.name.lower()=='mlp':
-        model = MLP()
-    elif config.model.name.lower()=='cnn':
-        model = CNN()
-    else:
-        raise NotImplementedError(f"Model {config.model.name} unknown.")
-    return model
+  if config.model.name.lower()=='mlp':
+    model = MLP()
+  elif config.model.name.lower()=='cnn':
+    model = CNN()
+  else:
+    raise NotImplementedError(f"Model {config.model.name} unknown.")
+  return model
 
 
 def get_solver(config, sde, score):
-    if config.solver.outer_solver.lower()=="eulermaruyama":
-        outer_solver = EulerMaruyama(sde.reverse(score),
-                                     num_steps=config.solver.num_outer_steps,
-                                     dt=config.solver.dt, epsilon=config.solver.epsilon)
-    else:
-        raise NotImplementedError(f"Solver {config.solver.outer_solver} unknown.")
-    if config.solver.inner_solver is None:
-        inner_solver = None
-    elif config.solver.inner_solver.lower()=="annealed":
-        inner_solver = Annealed(sde.corrector(sde_lib.UDLangevin, score), num_steps=config.solver.num_inner_steps, snr=config.solver.snr)
-    else:
-        raise NotImplementedError(f"Solver {config.solver.inner_solver} unknown.")
-    return outer_solver, inner_solver
+  if config.solver.outer_solver.lower()=="eulermaruyama":
+    outer_solver = EulerMaruyama(sde.reverse(score),
+                                 num_steps=config.solver.num_outer_steps,
+                                 dt=config.solver.dt, epsilon=config.solver.epsilon)
+  else:
+    raise NotImplementedError(f"Solver {config.solver.outer_solver} unknown.")
+  if config.solver.inner_solver is None:
+    inner_solver = None
+  elif config.solver.inner_solver.lower()=="annealed":
+    inner_solver = Annealed(sde.corrector(sde_lib.UDLangevin, score), num_steps=config.solver.num_inner_steps, snr=config.solver.snr)
+  else:
+    raise NotImplementedError(f"Solver {config.solver.inner_solver} unknown.")
+  return outer_solver, inner_solver
 
 
 def get_ddim_chain(config, model):
-    """
-    Args:
-        model: DDIM parameterizes the `\epsilon(x, t) = -1. * fwd_marginal_std(t) * score(x, t)` function
-    """
-    if config.solver.outer_solver.lower()=="ddimvp":
-        outer_solver = DDIMVP(model, eta=config.solver.eta, num_steps=config.solver.num_outer_steps,
-                            dt=config.solver.dt, epsilon=config.solver.epsilon,
-                            beta_min=config.model.beta_min,
-                            beta_max=config.model.beta_max)
-    elif config.solver.outer_solver.lower()=="ddimve":
-        outer_solver = DDIMVE(model, eta=config.solver.eta, num_steps=config.solver.num_outer_steps,
-                            dt=config.solver.dt, epsilon=config.solver.epsilon,
-                            sigma_min=config.model.sigma_min,
-                            sigma_max=config.model.sigma_max)
-    else:
-        raise NotImplementedError(f"DDIM Chain {config.solver.outer_solver} unknown.")
-    return outer_solver
+  """
+  Args:
+      model: DDIM parameterizes the `\epsilon(x, t) = -1. * fwd_marginal_std(t) * score(x, t)` function
+  """
+  if config.solver.outer_solver.lower()=="ddimvp":
+    outer_solver = DDIMVP(model, eta=config.solver.eta, num_steps=config.solver.num_outer_steps,
+                          dt=config.solver.dt, epsilon=config.solver.epsilon,
+                          beta_min=config.model.beta_min,
+                          beta_max=config.model.beta_max)
+  elif config.solver.outer_solver.lower()=="ddimve":
+    outer_solver = DDIMVE(model, eta=config.solver.eta, num_steps=config.solver.num_outer_steps,
+                          dt=config.solver.dt, epsilon=config.solver.epsilon,
+                          sigma_min=config.model.sigma_min,
+                          sigma_max=config.model.sigma_max)
+  else:
+    raise NotImplementedError(f"DDIM Chain {config.solver.outer_solver} unknown.")
+  return outer_solver
 
 
 def get_markov_chain(config, score):
-    """
-    Args:
-        score: DDPM/SMLD(NCSN) parameterizes the `score(x, t)` function.
-    """
-    if config.solver.outer_solver.lower()=="ddpm":
-        outer_solver = DDPM(score, num_steps=config.solver.num_outer_steps,
-                            dt=config.solver.dt, epsilon=config.solver.epsilon,
-                            beta_min=config.model.beta_min,
-                            beta_max=config.model.beta_max)
-    elif config.solver.outer_solver.lower()=="smld":
-        outer_solver = SMLD(score, num_steps=config.solver.num_outer_steps,
-                            dt=config.solver.dt, epsilon=config.solver.epsilon,
-                            sigma_min=config.model.sigma_min,
-                            sigma_max=config.model.sigma_max)
-    else:
-        raise NotImplementedError(f"Markov Chain {config.solver.outer_solver} unknown.")
-    return outer_solver
+  """
+  Args:
+    score: DDPM/SMLD(NCSN) parameterizes the `score(x, t)` function.
+  """
+  if config.solver.outer_solver.lower()=="ddpm":
+    outer_solver = DDPM(score, num_steps=config.solver.num_outer_steps,
+                        dt=config.solver.dt, epsilon=config.solver.epsilon,
+                        beta_min=config.model.beta_min,
+                        beta_max=config.model.beta_max)
+  elif config.solver.outer_solver.lower()=="smld":
+    outer_solver = SMLD(score, num_steps=config.solver.num_outer_steps,
+                        dt=config.solver.dt, epsilon=config.solver.epsilon,
+                        sigma_min=config.model.sigma_min,
+                        sigma_max=config.model.sigma_max)
+  else:
+    raise NotImplementedError(f"Markov Chain {config.solver.outer_solver} unknown.")
+  return outer_solver
 
 
 def numpy_collate(batch):
-    if isinstance(batch[0], np.ndarray):
-        return np.stack(batch)
-    elif isinstance(batch[0], (tuple, list)):
-        transposed = zip(*batch)
-        return [numpy_collate(samples) for samples in transposed]
-    else:
-        return np.array(batch)
+  if isinstance(batch[0], np.ndarray):
+    return np.stack(batch)
+  elif isinstance(batch[0], (tuple, list)):
+    transposed = zip(*batch)
+    return [numpy_collate(samples) for samples in transposed]
+  else:
+    return np.array(batch)
 
 
 def jit_collate(n_jitted_steps, batch_size, batch):
-    return np.reshape(batch, (n_jitted_steps, batch_size, -1))
+  return np.reshape(batch, (n_jitted_steps, batch_size, -1))
 
 
 def pmap_and_jit_collate(num_devices, n_jitted_steps, per_device_batch_size, batch):
-    return np.reshape(batch, (num_devices, n_jitted_steps, per_device_batch_size, -1))
+  return np.reshape(batch, (num_devices, n_jitted_steps, per_device_batch_size, -1))
 
 
 def pmap_collate(num_devices, per_device_batch_size, batch):
-    return np.reshape(batch, (num_devices, per_device_batch_size, -1))
+  return np.reshape(batch, (num_devices, per_device_batch_size, -1))
 
 
 class NumpyLoader(DataLoader):
-    def __init__(self, config, dataset,
-                 shuffle=False, sampler=None,
-                 batch_sampler=None, num_workers=0,
-                 pin_memory=False, drop_last=False,
-                 timeout=0, worker_init_fn=None):
-        prod_batch_dims = config.training.batch_size * config.training.n_jitted_steps
-        if config.training.pmap and config.training.n_jitted_steps != 1:
-            collate_fn = partial(
-                pmap_and_jit_collate, jax.local_device_count(),
-                config.training.n_jitted_steps,
-                config.training.batch_size // jax.local_device_count())
-        elif config.training.pmap and config.training.n_jitted_steps==1:
-            collate_fn = partial(
-                pmap_collate, jax.local_device_count(),
-                config.training.batch_size // jax.local_device_count())
-        elif config.training.n_jitted_steps != 1:
-            collate_fn = partial(
-                jit_collate, config.training.n_jitted_steps, config.training.batch_size)
-        else:
-            collate_fn = numpy_collate
+  def __init__(self, config, dataset,
+               shuffle=False, sampler=None,
+               batch_sampler=None, num_workers=0,
+               pin_memory=False, drop_last=False,
+               timeout=0, worker_init_fn=None):
+    prod_batch_dims = config.training.batch_size * config.training.n_jitted_steps
+    if config.training.pmap and config.training.n_jitted_steps != 1:
+      collate_fn = partial(
+        pmap_and_jit_collate, jax.local_device_count(),
+        config.training.n_jitted_steps,
+        config.training.batch_size // jax.local_device_count())
+    elif config.training.pmap and config.training.n_jitted_steps==1:
+      collate_fn = partial(
+        pmap_collate, jax.local_device_count(),
+        config.training.batch_size // jax.local_device_count())
+    elif config.training.n_jitted_steps != 1:
+      collate_fn = partial(
+        jit_collate, config.training.n_jitted_steps, config.training.batch_size)
+    else:
+      collate_fn = numpy_collate
 
-        super(self.__class__, self).__init__(dataset,
-                                             batch_size=prod_batch_dims,
-                                             shuffle=shuffle,
-                                             sampler=sampler,
-                                             batch_sampler=batch_sampler,
-                                             num_workers=num_workers,
-                                             collate_fn=collate_fn,
-                                             pin_memory=pin_memory,
-                                             drop_last=drop_last,
-                                             timeout=timeout,
-                                             worker_init_fn=worker_init_fn)
+    super(self.__class__, self).__init__(dataset,
+                                         batch_size=prod_batch_dims,
+                                         shuffle=shuffle,
+                                         sampler=sampler,
+                                         batch_sampler=batch_sampler,
+                                         num_workers=num_workers,
+                                         collate_fn=collate_fn,
+                                         pin_memory=pin_memory,
+                                         drop_last=drop_last,
+                                         timeout=timeout,
+                                         worker_init_fn=worker_init_fn)
 
 
 def train(sampling_shape, config, workdir, dataset):
-    """ Train a score based generative model using stochastic gradient descent
-
-    Args:
-        sampling_shape : sampling shape may differ depending on the modality of data
-        config: An ml-collections configuration to use.
-        workdir: Working directory for checkpoints and TF summaries. If this
-            contains checkpoint training will be resumed from the latest checkpoint.
-        dataset: a valid `torch.DataLoader` class.
-
-    Returns:
-
-    """
-    train_dataloader = NumpyLoader(config, dataset)
-    eval_dataloader = NumpyLoader(config, dataset)
-
-    jax.default_device = jax.devices()[0]
-    # Tip: use `export CUDA_VISIBLE_DEVICES` to restrict the devices visible to jax
-    # ... devices (GPUs/TPUs) must be all the same model for pmap to work
-    num_devices =  int(jax.local_device_count())
-    if jax.process_index()==0: print("num_devices={}, pmap={}".format(
-        num_devices, config.training.pmap))
-
-    # Create directories for experimental logs
-    sample_dir = os.path.join(workdir, "samples")
-    if not os.path.exists(sample_dir):
-        os.mkdir(sample_dir)
-
-    scaler = dataset.get_data_scaler(config)
-    inverse_scaler = dataset.get_data_inverse_scaler(config)
-
-    # eval_function = dataset.calculate_metrics_batch
-    # metric_names = dataset.metric_names()
-
-    if jax.process_index()==0:
-      run = wandb.init(
-          project="diffusionjax",
-          config=config,
-      )
+  """ Train a score based generative model using stochastic gradient descent
+
+  Args:
+    sampling_shape : sampling shape may differ depending on the modality of data
+    config: An ml-collections configuration to use.
+    workdir: Working directory for checkpoints and TF summaries. If this
+      contains checkpoint training will be resumed from the latest checkpoint.
+    dataset: a valid `torch.DataLoader` class.
+  """
+  train_dataloader = NumpyLoader(config, dataset)
+  eval_dataloader = NumpyLoader(config, dataset)
+
+  jax.default_device = jax.devices()[0]
+  # Tip: use `export CUDA_VISIBLE_DEVICES` to restrict the devices visible to jax
+  # ... devices (GPUs/TPUs) must be all the same model for pmap to work
+  num_devices =  int(jax.local_device_count())
+  if jax.process_index()==0: print("num_devices={}, pmap={}".format(
+    num_devices, config.training.pmap))
+
+  # Create directories for experimental logs
+  sample_dir = os.path.join(workdir, "samples")
+  if not os.path.exists(sample_dir):
+    os.mkdir(sample_dir)
+
+  scaler = dataset.get_data_scaler(config)
+  inverse_scaler = dataset.get_data_inverse_scaler(config)
+  # eval_function = dataset.calculate_metrics_batch
+  # metric_names = dataset.metric_names()
+  if jax.process_index()==0:
+    run = wandb.init(
+      project="diffusionjax",
+      config=config,
+    )
+  rng = random.PRNGKey(config.seed)
+
+  # Initialize model
+  rng, model_rng = random.split(rng, 2)
+  model = get_model(config)
+  # Initialize parameters
+  params = model.init(
+    model_rng,
+    jnp.zeros(
+      sampling_shape
+    ),
+    jnp.ones((sampling_shape[0],)))
+
+  # Initialize optimizer
+  optimizer = get_optimizer(config)
+  opt_state = optimizer.init(params)
+  state = State(step=0,
+    opt_state=opt_state,
+    params=params,
+    lr=config.optim.lr,
+    rng=rng)
+
+  # Create checkpoints directory
+  checkpoint_dir = os.path.join(workdir, "checkpoints")
+  if not os.path.exists(checkpoint_dir):
+    os.mkdir(checkpoint_dir)
+
+  # Intermediate checkpoints to resume training after pre-emption in cloud environments
+  checkpoint_meta_dir = os.path.join(workdir, "checkpoints-meta")
+  if not os.path.exists(checkpoint_meta_dir):
+    os.mkdir(checkpoint_meta_dir)
+
+  # Orbax checkpointer boilerplate
+  manager_options = orbax.checkpoint.CheckpointManagerOptions(
+    create=True, max_to_keep=np.inf)
+  checkpoint_manager = orbax.checkpoint.CheckpointManager(
+    checkpoint_dir,
+    orbax.checkpoint.Checkpointer(orbax.checkpoint.PyTreeCheckpointHandler()), manager_options)
+
+  # meta_manager_options = orbax.checkpoint.CheckpointManagerOptions(
+  #   create=True, max_to_keep=1)
+  meta_checkpoint_manager = orbax.checkpoint.CheckpointManager(
+    checkpoint_meta_dir,
+    orbax.checkpoint.Checkpointer(orbax.checkpoint.PyTreeCheckpointHandler()), manager_options)
+
+  # Resume training when intermediate checkpoints are detected
+  restore_args = orbax_utils.restore_args_from_target(state, mesh=None)
+  save_step = meta_checkpoint_manager.latest_step()
+  if save_step is not None:
+    meta_checkpoint_manager.restore(
+      save_step,
+      items=state, restore_kwargs={'restore_args': restore_args})
+
+  # `state.step` is JAX integer on the GPU/TPU devices
+  initial_step = int(state.step)
+  rng = state.rng
+
+  # Build one-step training and evaluation functions
+  sde = get_sde(config)
+  # Trained score
+  score = get_score(sde, model, params, score_scaling=config.training.score_scaling)
+
+  # Setup solver
+  outer_solver, inner_solver = get_solver(config, sde, score)
+
+  loss = get_loss(
+    sde, outer_solver, model,
+    score_scaling=config.training.score_scaling,
+    likelihood_weighting=config.training.likelihood_weighting)
+  train_step = get_step_fn(loss, optimizer, train=True, pmap=config.training.pmap)
+  eval_step = get_step_fn(loss, optimizer, train=False, pmap=config.training.pmap)
+
+  if config.training.n_jitted_steps > 1:
+    train_step = partial(jax.lax.scan, train_step)
+    eval_step = partial(jax.lax.scan, eval_step)
+
+  if config.training.pmap:
+    train_step = jax.pmap(
+      train_step, axis_name='batch', donate_argnums=1)
+    eval_step = jax.pmap(
+      eval_step, axis_name='batch', donate_argnums=1)
+
+  # Replicate the training state to run on multiple devices
+  if config.training.pmap: state = flax_utils.replicate(state)
+
+  # Probably want to train over multiple epochs
+  # If num_epochs > num_batch, decides which tqdm to go over
+  i_epoch = 0
+  prev_time = time.time()
+
+  # Deal with training in a number of steps
+  # num_epochs = config.training.n_iters // (dataset_size / batch_size)
+  num_epochs = config.training.n_iters
+
+  step = initial_step
+  # In case there are multiple hosts (e.g., TPU pods), only log to host 0
+  if jax.process_index() == 0:
+    logging.info("Starting training loop at step %d." % (initial_step,))
+  rng = jax.random.fold_in(rng, jax.process_index())
+
+  # JIT multiple training steps together for faster training
+  n_jitted_steps = config.training.n_jitted_steps
+
+  # Must be divisible by the number of steps jitted together
+  assert config.training.log_step_freq % n_jitted_steps == 0 and \
+    config.training.snapshot_freq_for_preemption % n_jitted_steps == 0 and \
+    config.training.eval_freq % n_jitted_steps == 0 and \
+    config.training.snapshot_freq % n_jitted_steps == 0, "Missing logs or checkpoints!"
+
+  mean_losses = jnp.zeros((num_epochs, 1))
+  for i_epoch in trange(1, num_epochs + 1, unit="epochs"):
+    current_time = time.time()
+
+    if i_epoch != 0 and (num_epochs < config.training.batch_size):
+      print("Epoch took {:.1f} seconds".format(current_time - prev_time))
+      prev_time = time.time()
+
+    eval_iter = iter(eval_dataloader)
+
+    with tqdm(
+      train_dataloader,
+      unit=" batch",
+      disable=True
+    ) as tepoch:
+      tepoch.set_description(f"Epoch {i_epoch}")
+      losses = jnp.empty((len(tepoch), 1))
+
+      for i_batch, batch in enumerate(tepoch):
+        batch = jax.tree_map(lambda x: scaler(x), batch)
+
+        # Execute one training step
+        if config.training.pmap:
+          rng, *next_rng = jax.random.split(rng, num=jax.local_device_count() + 1)
+          next_rng = jnp.asarray(next_rng)
+        else:
+          rng, next_rng = jax.random.split(rng, num=2)
 
-    rng = random.PRNGKey(config.seed)
-
-    # Initialize model
-    rng, model_rng = random.split(rng, 2)
-    model = get_model(config)
-    # Initialize parameters
-    params = model.init(
-        model_rng,
-        jnp.zeros(
-            sampling_shape
-        ),
-        jnp.ones((sampling_shape[0],)))
-
-    # Initialize optimizer
-    optimizer = get_optimizer(config)
-    opt_state = optimizer.init(params)
-
-    state = State(step=0,
-        opt_state=opt_state,
-        params=params,
-        lr=config.optim.lr,
-        rng=rng)
-
-    # Create checkpoints directory
-    checkpoint_dir = os.path.join(workdir, "checkpoints")
-    if not os.path.exists(checkpoint_dir):
-        os.mkdir(checkpoint_dir)
-
-    # Intermediate checkpoints to resume training after pre-emption in cloud environments
-    checkpoint_meta_dir = os.path.join(workdir, "checkpoints-meta")
-    if not os.path.exists(checkpoint_meta_dir):
-        os.mkdir(checkpoint_meta_dir)
-
-    # Orbax checkpointer boilerplate
-    manager_options = orbax.checkpoint.CheckpointManagerOptions(
-        create=True, max_to_keep=np.inf)
-    checkpoint_manager = orbax.checkpoint.CheckpointManager(
-        checkpoint_dir,
-        orbax.checkpoint.Checkpointer(orbax.checkpoint.PyTreeCheckpointHandler()), manager_options)
-
-    meta_manager_options = orbax.checkpoint.CheckpointManagerOptions(
-        create=True, max_to_keep=1)
-    meta_checkpoint_manager = orbax.checkpoint.CheckpointManager(
-        checkpoint_meta_dir,
-        orbax.checkpoint.Checkpointer(orbax.checkpoint.PyTreeCheckpointHandler()), manager_options)
-
-    # Resume training when intermediate checkpoints are detected
-    restore_args = flax.training.orbax_utils.restore_args_from_target(state, mesh=None)
-    save_step = meta_checkpoint_manager.latest_step()
-    if save_step is not None:
-        meta_checkpoint_manager.restore(
-            save_step,
-            items=state, restore_kwargs={'restore_args': restore_args})
-
-    # `state.step` is JAX integer on the GPU/TPU devices
-    initial_step = int(state.step)
-    rng = state.rng
-
-    # Build one-step training and evaluation functions
-    sde = get_sde(config)
-    # Trained score
-    score = get_score(sde, model, params, score_scaling=config.training.score_scaling)
-
-    # Setup solver
-    outer_solver, inner_solver = get_solver(config, sde, score)
-
-    loss = get_loss(
-        sde, outer_solver, model,
-        score_scaling=config.training.score_scaling,
-        likelihood_weighting=config.training.likelihood_weighting)
-
-    train_step = get_step_fn(loss, optimizer, train=True, pmap=config.training.pmap)
-    eval_step = get_step_fn(loss, optimizer, train=False, pmap=config.training.pmap)
-
-    if config.training.n_jitted_steps > 1:
-        train_step = partial(jax.lax.scan, train_step)
-        eval_step = partial(jax.lax.scan, eval_step)
-
-    if config.training.pmap:
-        train_step = jax.pmap(
-            train_step, axis_name='batch', donate_argnums=1)
-        eval_step = jax.pmap(
-            eval_step, axis_name='batch', donate_argnums=1)
-
-    # Replicate the training state to run on multiple devices
-    if config.training.pmap: state = flax_utils.replicate(state)
-
-    # Probably want to train over multiple epochs
-    # If num_epochs > num_batch, decides which tqdm to go over
-    i_epoch = 0
-    prev_time = time.time()
-
-    # Deal with training in a number of steps
-    # num_epochs = config.training.n_iters // (dataset_size / batch_size)
-    num_epochs = config.training.n_iters
-
-    step = initial_step
-    # In case there are multiple hosts (e.g., TPU pods), only log to host 0
-    if jax.process_index() == 0:
-        logging.info("Starting training loop at step %d." % (initial_step,))
-    rng = jax.random.fold_in(rng, jax.process_index())
-
-    # JIT multiple training steps together for faster training
-    n_jitted_steps = config.training.n_jitted_steps
-
-    # Must be divisible by the number of steps jitted together
-    assert config.training.log_step_freq % n_jitted_steps == 0 and \
-        config.training.snapshot_freq_for_preemption % n_jitted_steps == 0 and \
-        config.training.eval_freq % n_jitted_steps == 0 and \
-        config.training.snapshot_freq % n_jitted_steps == 0, "Missing logs or checkpoints!"
-
-    mean_losses = jnp.zeros((num_epochs, 1))
-    for i_epoch in trange(1, num_epochs + 1, unit="epochs"):
-        current_time = time.time()
-
-        if i_epoch != 0 and (num_epochs < config.training.batch_size):
-            print("Epoch took {:.1f} seconds".format(current_time - prev_time))
-            prev_time = time.time()
-
-        eval_iter = iter(eval_dataloader)
-
-        with tqdm(
-                train_dataloader,
-                unit=" batch",
-                disable=True
-        ) as tepoch:
-            tepoch.set_description(f"Epoch {i_epoch}")
-            losses = jnp.empty((len(tepoch), 1))
-
-            for i_batch, batch in enumerate(tepoch):
-                batch = jax.tree_map(lambda x: scaler(x), batch)
-
-                # Execute one training step
-                if config.training.pmap:
-                    rng, *next_rng = jax.random.split(rng, num=jax.local_device_count() + 1)
-                    next_rng = jnp.asarray(next_rng)
-                else:
-                    rng, next_rng = jax.random.split(rng, num=2)
-
-                (_, params, opt_state), loss_train = train_step(
-                    (next_rng, state.params, state.opt_state), batch)
-                state = state.replace(opt_state=opt_state, params=params)
-
-                if config.training.pmap:
-                    loss_train = flax_utils.unreplicate(loss_train).mean()  # returns a single instance of replicated loss array
-                else:
-                    loss_train = loss_train.mean()
-
-                # Log to console, file and wandb on host 0
-                if jax.process_index()==0:
-                    step += config.training.n_jitted_steps
-                    losses = losses.at[i_batch].set(loss_train)
-                    if step % config.training.log_step_freq==0 and jax.process_index() == 0:
-                        logging.info("step {:d}, training_loss {:.2e}".format(step, loss_train))
-
-                # Save a temporary checkpoint to resume training after pre-emption (for cloud computing environments) periodically
-                if step!=0 and step % config.training.snapshot_freq_for_preemption==0 and jax.process_index()==0:
-                    if config.training.pmap:
-                        saved_state = flax_utils.unreplicate(state)
-                    else:
-                        saved_state = state
-                    saved_state = saved_state.replace(rng=rng)
-                    saved_args = flax.training.orbax_utils.save_args_from_target(saved_state)
-                    meta_checkpoint_manager.save(step//config.training.snapshot_freq_for_preemption, saved_state, save_kwargs={'save_args': saved_args})
-
-                # Report the loss on an evaluation dataset periodically
-                if step % config.training.eval_freq == 0:
-                    eval_batch = jax.tree_map(lambda x: scaler(x), next(eval_iter))
-                    if config.training.pmap:
-                        rng, *next_rng = jax.random.split(rng, num=jax.local_device_count() + 1)
-                        next_rng = jnp.asarray(next_rng)
-                    else:
-                        rng, next_rng = jax.random.split(rng, num=2)
-                    (_, _, _), loss_eval = eval_step(
-                        (next_rng, state.params, state.opt_state), eval_batch)
-
-                    if config.training.pmap:
-                        loss_eval = flax_utils.unreplicate(loss_eval).mean()
-                    else:
-                        loss_eval = loss_eval.mean()
-
-                    if jax.process_index()==0:
-                        logging.info("batch: {:d}, eval_loss: {:.5e}".format(step, loss_eval))
-                        wandb.log({"eval-loss": loss_eval})
-
-                # Save a checkpoint periodically and generate samples if needed
-                if step != 0 and step % config.training.snapshot_freq == 0 or step == config.training.n_iters:
-                    # Save the checkpoint
-                    if jax.process_index()==0:
-                        if config.training.pmap:
-                            saved_state = flax_utils.unreplicate(state)
-                        else:
-                            saved_state = state
-                        saved_state = saved_state.replace(rng=rng)
-                        saved_args = flax.training.orbax_utils.save_args_from_target(saved_state)
-                        checkpoint_manager.save(step // config.training.snapshot_freq, saved_state, save_kwargs={'save_args': saved_args})
-
-                    # Generate and save samples
-                    if config.training.snapshot_sampling:
-                        # Setup solver with new trained score
-                        # Use the unreplicated parameters of the saved state
-                        score = get_score(
-                            sde, model, saved_state.params, config.training.score_scaling)
-                        outer_solver, inner_solver = get_solver(config, sde, score)
-                        sampler = get_sampler(sampling_shape, outer_solver,
-                                            inner_solver, denoise=config.sampling.denoise,
-                                            stack_samples=config.sampling.stack_samples,
-                                            inverse_scaler=inverse_scaler)
-
-                        if config.training.pmap:
-                            sampler = jax.pmap(sampler, axis_name='batch')
-                            rng, *sample_rng = random.split(rng, 1 + jax.local_device_count())
-                            sample_rng = jnp.asarray(sample_rng)
-                        else:
-                            rng, sample_rng = random.split(rng, 2)
-
-                        sample, nfe = sampler(sample_rng)
-
-                        # eval_fn = eval_function(sample)
-                        # wandb.log({metric_names[0]: eval_fn})
-
-                        this_sample_dir = os.path.join(
-                            sample_dir, "iter_{}_host_{}".format(step, jax.process_index()))
-                        if not os.path.isdir(this_sample_dir):
-                            os.mkdir(this_sample_dir)
-
-                        with open(os.path.join(this_sample_dir, "sample.np"), 'wb') as infile:
-                            np.save(infile, sample)
-
-            if jax.process_index()==0:
-                mean_loss = jnp.mean(losses, axis=0)
-
-        if jax.process_index()==0 and i_epoch % config.training.log_epoch_freq==0:
-            logging.info("step {:d}, mean_loss {:.2e}".format(step, mean_loss[0]))
-            mean_losses = mean_losses.at[i_epoch].set(mean_loss)
-            wandb.log({"train-loss": mean_loss})
-
-    artifact = wandb.Artifact(name='checkpoint', type='checkpoint')
-    artifact.add_dir(local_path=checkpoint_dir)
-    run.log_artifact(artifact)
-
-    # Get the model and do test dataset
-    if config.training.pmap:
-        saved_state = flax_utils.unreplicate(state)
-    else:
-        saved_state = state
-    return saved_state.params, saved_state.opt_state, mean_losses
+        (_, params, opt_state), loss_train = train_step(
+          (next_rng, state.params, state.opt_state), batch)
+        state = state.replace(opt_state=opt_state, params=params)
+
+        if config.training.pmap:
+          loss_train = flax_utils.unreplicate(loss_train).mean()  # returns a single instance of replicated loss array
+        else:
+          loss_train = loss_train.mean()
+
+        # Log to console, file and wandb on host 0
+        if jax.process_index()==0:
+          step += config.training.n_jitted_steps
+          losses = losses.at[i_batch].set(loss_train)
+          if step % config.training.log_step_freq==0 and jax.process_index() == 0:
+            logging.info("step {:d}, training_loss {:.2e}".format(step, loss_train))
+
+        # Save a temporary checkpoint to resume training after pre-emption (for cloud computing environments) periodically
+        if step!=0 and step % config.training.snapshot_freq_for_preemption==0 and jax.process_index()==0:
+          if config.training.pmap:
+            saved_state = flax_utils.unreplicate(state)
+          else:
+            saved_state = state
+          saved_state = saved_state.replace(rng=rng)
+          saved_args = orbax_utils.save_args_from_target(saved_state)
+          meta_checkpoint_manager.save(step//config.training.snapshot_freq_for_preemption, saved_state, save_kwargs={'save_args': saved_args})
+
+        # Report the loss on an evaluation dataset periodically
+        if step % config.training.eval_freq == 0:
+          eval_batch = jax.tree_map(lambda x: scaler(x), next(eval_iter))
+          if config.training.pmap:
+            rng, *next_rng = jax.random.split(rng, num=jax.local_device_count() + 1)
+            next_rng = jnp.asarray(next_rng)
+          else:
+            rng, next_rng = jax.random.split(rng, num=2)
+          (_, _, _), loss_eval = eval_step(
+            (next_rng, state.params, state.opt_state), eval_batch)
+
+          if config.training.pmap:
+            loss_eval = flax_utils.unreplicate(loss_eval).mean()
+          else:
+            loss_eval = loss_eval.mean()
+
+          if jax.process_index()==0:
+            logging.info("batch: {:d}, eval_loss: {:.5e}".format(step, loss_eval))
+            wandb.log({"eval-loss": loss_eval})
+
+        # Save a checkpoint periodically and generate samples if needed
+        if step != 0 and step % config.training.snapshot_freq == 0 or step == config.training.n_iters:
+          # Save the checkpoint
+          if jax.process_index()==0:
+            if config.training.pmap:
+              saved_state = flax_utils.unreplicate(state)
+            else:
+              saved_state = state
+            saved_state = saved_state.replace(rng=rng)
+            saved_args = orbax_utils.save_args_from_target(saved_state)
+            checkpoint_manager.save(step // config.training.snapshot_freq, saved_state, save_kwargs={'save_args': saved_args})
+
+          # Generate and save samples
+          if config.training.snapshot_sampling:
+            # Setup solver with new trained score
+            # Use the unreplicated parameters of the saved state
+            score = get_score(
+              sde, model, saved_state.params, config.training.score_scaling)
+            outer_solver, inner_solver = get_solver(config, sde, score)
+            sampler = get_sampler(sampling_shape, outer_solver,
+                                  inner_solver, denoise=config.sampling.denoise,
+                                  stack_samples=config.sampling.stack_samples,
+                                  inverse_scaler=inverse_scaler)
+
+            if config.training.pmap:
+              sampler = jax.pmap(sampler, axis_name='batch')
+              rng, *sample_rng = random.split(rng, 1 + jax.local_device_count())
+              sample_rng = jnp.asarray(sample_rng)
+            else:
+              rng, sample_rng = random.split(rng, 2)
+
+            sample, _ = sampler(sample_rng)
+
+            # eval_fn = eval_function(sample)
+            # wandb.log({metric_names[0]: eval_fn})
+
+            this_sample_dir = os.path.join(
+              sample_dir, "iter_{}_host_{}".format(step, jax.process_index()))
+            if not os.path.isdir(this_sample_dir):
+              os.mkdir(this_sample_dir)
+
+            with open(os.path.join(this_sample_dir, "sample.np"), 'wb') as infile:
+              np.save(infile, sample)
+
+      if jax.process_index()==0:
+        mean_loss = jnp.mean(losses, axis=0)
+
+    if jax.process_index()==0 and i_epoch % config.training.log_epoch_freq==0:
+      logging.info("step {:d}, mean_loss {:.2e}".format(step, mean_loss[0]))
+      mean_losses = mean_losses.at[i_epoch].set(mean_loss)
+      wandb.log({"train-loss": mean_loss})
+
+  artifact = wandb.Artifact(name='checkpoint', type='checkpoint')
+  artifact.add_dir(local_path=checkpoint_dir)
+  run.log_artifact(artifact)
+
+  # Get the model and do test dataset
+  if config.training.pmap:
+    saved_state = flax_utils.unreplicate(state)
+  else:
+    saved_state = state
+  return saved_state.params, saved_state.opt_state, mean_losses
diff --git a/diffusionjax/sde.py b/diffusionjax/sde.py
index 179dd30..60e00d2 100644
--- a/diffusionjax/sde.py
+++ b/diffusionjax/sde.py
@@ -1,182 +1,190 @@
 """SDE class."""
 import abc
-from functools import partial
-import jax
 import jax.numpy as jnp
 from jax import random
 from diffusionjax.utils import batch_mul
 
 
 class SDE(abc.ABC):
-    """SDE abstract class. Functions are designed for a mini-batch of inputs."""
+  """SDE abstract class. Functions are designed for a mini-batch of inputs."""
 
-    def __init__(self):
-        """Construct an SDE.
-        """
+  def __init__(self):
+    """Construct an SDE."""
 
-    @abc.abstractmethod
-    def sde(self, x, t):
-        r"""Return the drift and diffusion coefficients of the SDE.
+  @abc.abstractmethod
+  def sde(self, x, t):
+    """Return the drift and diffusion coefficients of the SDE.
 
-        Args:
-            x: a JAX tensor of the state
-            t: JAX float of the time
+    Args:
+      x: a JAX tensor of the state
+      t: JAX float of the time
 
-        Returns:
-            drift: drift function of the forward SDE
-            diffusion: dispersion function of the forward SDE
-        """
+    Returns:
+      drift: drift function of the forward SDE
+      diffusion: dispersion function of the forward SDE
+    """
 
-    class RSDE:
-        """Reverse SDE class."""
-        def __init__(self, score, forward_sde):
-            self.score = score
-            self.forward_sde = forward_sde
+  class RSDE:
+    """Reverse SDE class."""
+    def __init__(self, score, forward_sde):
+      self.score = score
+      self.forward_sde = forward_sde
 
-        def sde(self, x, t):
-            drift, diffusion = self.forward_sde(x, t)
-            drift = -drift + batch_mul(diffusion**2, self.score(x, t))
-            return drift, diffusion
+    def sde(self, x, t):
+      drift, diffusion = self.forward_sde(x, t)
+      drift = -drift + batch_mul(diffusion**2, self.score(x, t))
+      return drift, diffusion
 
-    def corrector(self, Corrector, score):
+  def corrector(self, Corrector, score):
 
-        class CSDE(Corrector, self.__class__):
-            def __init__(self):
-                super().__init__(score)
+    class CSDE(Corrector, self.__class__):
+      def __init__(self):
+        super().__init__(score)
 
-        return CSDE()
+    return CSDE()
 
 
 class ODLangevin(SDE):
-    """Overdamped langevin SDE."""
-    def __init__(self, score, damping=2e0, L=1.0):
-        super().__init__()
-        self.score = score
-        self.damping = damping
-        self.L = L
+  """Overdamped langevin SDE."""
+  def __init__(self, score, damping=2e0, L=1.0):
+    super().__init__()
+    self.score = score
+    self.damping = damping
+    self.L = L
 
-    def sde(self, x, t):
-        drift = -self.score(x, t)
-        diffusion = jnp.ones(x.shape) * jnp.sqrt(2 * self.damping / self.L)
-        return drift, diffusion
+  def sde(self, x, t):
+    drift = -self.score(x, t)
+    diffusion = jnp.ones(x.shape) * jnp.sqrt(2 * self.damping / self.L)
+    return drift, diffusion
 
 
 class UDLangevin(SDE):
-    """Underdamped Langevin SDE."""
-    def __init__(self, score):
-        super().__init__()
-        self.score = score
+  """Underdamped Langevin SDE."""
+  def __init__(self, score):
+    super().__init__()
+    self.score = score
 
-    def sde(self, x, t):
-        drift = -self.score(x, t)
-        diffusion = jnp.ones(x.shape) * jnp.sqrt(2)
-        return drift, diffusion
+  def sde(self, x, t):
+    drift = -self.score(x, t)
+    diffusion = jnp.ones(x.shape) * jnp.sqrt(2)
+    return drift, diffusion
 
 
 class VE(SDE):
-    """Variance exploding (VE) SDE, also known as diffusion process
-    with a time dependent diffusion coefficient.
+  """Variance exploding (VE) SDE, also known as diffusion process
+  with a time dependent diffusion coefficient.
+  """
+  def __init__(self, sigma_min=0.1, sigma_max=20.0):
+    super().__init__()
+    self.sigma_min = sigma_min
+    self.sigma_max = sigma_max
+
+  def sde(self, x, t):
+    sigma = self.sigma_min * (self.sigma_max / self.sigma_min)**t
+    drift = jnp.zeros_like(x)
+    diffusion = sigma * jnp.sqrt(2 * (jnp.log(self.sigma_max) - jnp.log(self.sigma_min)))
+    return drift, diffusion
+
+  def log_mean_coeff(self, t):
+    return jnp.zeros_like(t)
+
+  def mean_coeff(self, t):
+    return jnp.ones_like(t)
+
+  def variance(self, t):
+    std = self.sigma_min * (self.sigma_max / self.sigma_min)**t
+    return std**2
+
+  def marginal_prob(self, x, t):
+    r"""Parameters to determine the marginal distribution of the SDE,
+
+    .. math::
+      p_t(x)
     """
-    def __init__(self, sigma_min=0.1, sigma_max=20.0):
-        super().__init__()
-        self.sigma_min = sigma_min
-        self.sigma_max = sigma_max
+    std = self.sigma_min * (self.sigma_max / self.sigma_min)**t
+    mean = x
+    return mean, std
 
-    def sde(self, x, t):
-        sigma = self.sigma_min * (self.sigma_max / self.sigma_min)**t
-        drift = jnp.zeros_like(x)
-        diffusion = sigma * jnp.sqrt(2 * (jnp.log(self.sigma_max) - jnp.log(self.sigma_min)))
-        return drift, diffusion
+  def prior(self, rng, shape):
+    return random.normal(rng, shape) * self.sigma_max
 
-    def log_mean_coeff(self, t):
-        return jnp.zeros_like(t)
+  def reverse(self, score):
+    fwd_sde = self.sde
+    sigma_min = self.sigma_min
+    sigma_max = self.sigma_max
 
-    def mean_coeff(self, t):
-        return jnp.ones_like(t)
+    class RVE(self.RSDE, self.__class__):
+      def __init__(self):
+        super().__init__(score, fwd_sde)
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
 
-    def variance(self, t):
-        std = self.sigma_min * (self.sigma_max / self.sigma_min)**t
-        return std**2
+      def get_estimate_x_0(self, observation_map):
+        def estimate_x_0(x, t):
+          v_t = self.variance(t)
+          s = self.score(x, t)
+          x_0 = x + v_t * s
+          return observation_map(x_0), (s, x_0)
+        return estimate_x_0
 
-    def marginal_prob(self, x, t):
-        r"""Parameters to determine the marginal distribution of the SDE,
+    return RVE()
 
-        .. math::
-            p_t(x)
 
-        Args:
-            x: a JAX tensor of the state
-            t: JAX float of the time
-        """
-        std = self.sigma_min * (self.sigma_max / self.sigma_min)**t
-        mean = x
-        return mean, std
+class VP(SDE):
+  """Variance preserving (VP) SDE, also known
+  as time rescaled Ohrnstein Uhlenbeck (OU) SDE."""
+  def __init__(self, beta_min=0.1, beta_max=20.0):
+    super().__init__()
+    self.beta_min = beta_min
+    self.beta_max = beta_max
 
-    def prior(self, rng, shape):
-        return random.normal(rng, shape) * self.sigma_max
+  def sde(self, x, t):
+    beta_t = self.beta_min + t * (self.beta_max - self.beta_min)
+    drift = -0.5 * batch_mul(beta_t, x)
+    diffusion = jnp.sqrt(beta_t)
+    return drift, diffusion
 
-    def reverse(self, score):
-        fwd_sde = self.sde
-        sigma_min = self.sigma_min
-        sigma_max = self.sigma_max
+  def log_mean_coeff(self, t):
+    return -0.5 * t * self.beta_min - 0.25 * t**2 * (self.beta_max - self.beta_min)
 
-        class RVE(self.RSDE, self.__class__):
-            def __init__(self):
-                super().__init__(score, fwd_sde)
-                self.sigma_min = sigma_min
-                self.sigma_max = sigma_max
-        return RVE()
+  def mean_coeff(self, t):
+    return jnp.exp(self.log_mean_coeff(t))
 
+  def variance(self, t):
+    return 1.0 - jnp.exp(2 * self.log_mean_coeff(t))
 
-class VP(SDE):
-    """Variance preserving (VP) SDE, also known
-    as time rescaled Ohrnstein Uhlenbeck (OU) SDE."""
-    def __init__(self, beta_min=0.1, beta_max=20.0):
-        super().__init__()
+  def marginal_prob(self, x, t):
+    r"""Parameters to determine the marginal distribution of the SDE,
+
+    .. math::
+      p_t(x)
+    """
+    m = self.mean_coeff(t)
+    mean = batch_mul(m, x)
+    std = jnp.sqrt(self.variance(t))
+    return mean, std
+
+  def prior(self, rng, shape):
+    return random.normal(rng, shape)
+
+  def reverse(self, score):
+    fwd_sde = self.sde
+    beta_min = self.beta_min
+    beta_max = self.beta_max
+
+    class RVP(self.RSDE, self.__class__):
+      def __init__(self):
+        super().__init__(score, fwd_sde)
         self.beta_min = beta_min
         self.beta_max = beta_max
 
-    def sde(self, x, t):
-        beta_t = self.beta_min + t * (self.beta_max - self.beta_min)
-        drift = -0.5 * batch_mul(beta_t, x)
-        diffusion = jnp.sqrt(beta_t)
-        return drift, diffusion
-
-    def log_mean_coeff(self, t):
-        return -0.5 * t * self.beta_min - 0.25 * t**2 * (self.beta_max - self.beta_min)
-
-    def mean_coeff(self, t):
-        return jnp.exp(self.log_mean_coeff(t))
-
-    def variance(self, t):
-        return 1.0 - jnp.exp(2 * self.log_mean_coeff(t))
-
-    def marginal_prob(self, x, t):
-        r"""Parameters to determine the marginal distribution of the SDE,
-
-        .. math::
-            p_t(x)
-
-        Args:
-            x: a JAX tensor of the state
-            t: JAX float of the time
-        """
-        m = self.mean_coeff(t)
-        mean = batch_mul(m, x)
-        std = jnp.sqrt(self.variance(t))
-        return mean, std
-
-    def prior(self, rng, shape):
-        return random.normal(rng, shape)
-
-    def reverse(self, score):
-        fwd_sde = self.sde
-        beta_min = self.beta_min
-        beta_max = self.beta_max
-
-        class RVP(self.RSDE, self.__class__):
-            def __init__(self):
-                super().__init__(score, fwd_sde)
-                self.beta_min = beta_min
-                self.beta_max = beta_max
-        return RVP()
+      def get_estimate_x_0(self, observation_map):
+        def estimate_x_0(x, t):
+          m_t = self.mean_coeff(t)
+          v_t = self.variance(t)
+          s = self.score(x, t)
+          x_0 = (x + v_t * s) / m_t
+          return observation_map(x_0), (s, x_0)
+        return estimate_x_0
+
+    return RVP()
diff --git a/diffusionjax/solvers.py b/diffusionjax/solvers.py
index 3cc5a75..24aedc7 100644
--- a/diffusionjax/solvers.py
+++ b/diffusionjax/solvers.py
@@ -7,291 +7,350 @@
 
 
 class Solver(abc.ABC):
-    """Solver abstract class. Functions are designed for a mini-batch of inputs."""
-
-    def __init__(self, num_steps=1000, dt=None, epsilon=None):
-        """Construct a Solver.
-        Args:
-            num_steps: number of discretization time steps.
-            dt: time step duration, float or `None`.
-                Optional, if provided then final time, t1 = dt * num_steps.
-            epsilon: A small float 0. < `epsilon` << 1. The SDE or ODE are integrated to `epsilon` to avoid numerical issues.
-        """
-        self.num_steps = num_steps
-        if dt is not None:
-            self.t1 = dt * num_steps
-            if epsilon is not None:
-                # Defined in forward time, t \in [epsilon, t1], 0 < epsilon << t1
-                ts, step = jnp.linspace(epsilon, self.t1, num_steps, retstep=True)
-                self.ts = ts.reshape(-1, 1)
-                assert step == (self.t1 - epsilon) / num_steps
-                self.dt = step
-            else:
-                # Defined in forward time, t \in [dt , t1], 0 < \epsilon << t1
-                step = jnp.linspace(0, self.t1, num_steps + 1)
-                self.ts = ts[1:].reshape(-1, 1)
-                assert step == dt
-                self.dt = step
-        else:
-            self.t1 = 1.0
-            if epsilon is not None:
-                self.ts, step = jnp.linspace(epsilon, 1, num_steps, retstep=True)
-                assert step == (1. - epsilon) / num_steps
-                self.dt = step
-            else:
-                # Defined in forward time, t \in [dt, 1.0], 0 < dt << 1
-                ts, step = jnp.linspace(0, 1, num_steps + 1, retstep=True)
-                self.ts = ts[1:].reshape(-1, 1)
-                assert step == 1. / num_steps
-                self.dt = step
-
-    @abc.abstractmethod
-    def update(self, rng, x, t):
-        r"""Return the update of the state and any auxilliary values.
-
-        Args:
-            rng: A JAX random state.
-            x: A JAX array of the state.
-            t: JAX array of the time.
-
-        Returns:
-            x: A JAX array of the next state.
-            x_mean: A JAX array. The next state without random noise. Useful for denoising.
-        """
+  """Solver abstract class. Functions are designed for a mini-batch of inputs."""
+
+  def __init__(self, num_steps=1000, dt=None, epsilon=None):
+    """Construct a Solver.
+    Args:
+      num_steps: number of discretization time steps.
+      dt: time step duration, float or `None`.
+        Optional, if provided then final time, t1 = dt * num_steps.
+      epsilon: A small float 0. < `epsilon` << 1. The SDE or ODE are integrated to `epsilon` to avoid numerical issues.
+    """
+    self.num_steps = num_steps
+    if dt is not None:
+      self.t1 = dt * num_steps
+      if epsilon is not None:
+        # Defined in forward time, t \in [epsilon, t1], 0 < epsilon << t1
+        ts, step = jnp.linspace(epsilon, self.t1, num_steps, retstep=True)
+        self.ts = ts.reshape(-1, 1)
+        assert step == (self.t1 - epsilon) / num_steps
+        self.dt = step
+      else:
+        # Defined in forward time, t \in [dt , t1], 0 < \epsilon << t1
+        step = jnp.linspace(0, self.t1, num_steps + 1)
+        self.ts = ts[1:].reshape(-1, 1)
+        assert step == dt
+        self.dt = step
+    else:
+      self.t1 = 1.0
+      if epsilon is not None:
+        self.ts, step = jnp.linspace(epsilon, 1, num_steps, retstep=True)
+        assert step == (1. - epsilon) / num_steps
+        self.dt = step
+      else:
+        # Defined in forward time, t \in [dt, 1.0], 0 < dt << 1
+        ts, step = jnp.linspace(0, 1, num_steps + 1, retstep=True)
+        self.ts = ts[1:].reshape(-1, 1)
+        assert step == 1. / num_steps
+        self.dt = step
+
+  @abc.abstractmethod
+  def update(self, rng, x, t):
+    r"""Return the update of the state and any auxilliary values.
+
+    Args:
+      rng: A JAX random state.
+      x: A JAX array of the state.
+      t: JAX array of the time.
+
+    Returns:
+      x: A JAX array of the next state.
+      x_mean: A JAX array. The next state without random noise. Useful for denoising.
+    """
 
 
 class EulerMaruyama(Solver):
-    """Euler Maruyama numerical solver of an SDE.
-    Functions are designed for a mini-batch of inputs."""
-
-    def __init__(self, sde, num_steps=1000, dt=None, epsilon=None):
-        """Constructs an Euler-Maruyama Solver.
-        Args:
-            sde: A valid SDE class.
-        """
-        super().__init__(num_steps=num_steps, dt=dt, epsilon=epsilon)
-        self.sde = sde
-        self.prior = sde.prior
-
-    def update(self, rng, x, t):
-        drift, diffusion = self.sde.sde(x, t)
-        f = drift * self.dt
-        G = diffusion * jnp.sqrt(self.dt)
-        noise = random.normal(rng, x.shape)
-        x_mean = x + f
-        x = x_mean + batch_mul(G, noise)
-        return x, x_mean
+  """Euler Maruyama numerical solver of an SDE.
+  Functions are designed for a mini-batch of inputs."""
+
+  def __init__(self, sde, num_steps=1000, dt=None, epsilon=None):
+    """Constructs an Euler-Maruyama Solver.
+    Args:
+      sde: A valid SDE class.
+    """
+    super().__init__(num_steps=num_steps, dt=dt, epsilon=epsilon)
+    self.sde = sde
+    self.prior = sde.prior
+
+  def update(self, rng, x, t):
+    drift, diffusion = self.sde.sde(x, t)
+    f = drift * self.dt
+    G = diffusion * jnp.sqrt(self.dt)
+    noise = random.normal(rng, x.shape)
+    x_mean = x + f
+    x = x_mean + batch_mul(G, noise)
+    return x, x_mean
 
 
 class Annealed(Solver):
-    """Annealed Langevin numerical solver of an SDE.
-    Functions are designed for a mini-batch of inputs."""
-
-    def __init__(self, sde, snr=1e-2, num_steps=2, dt=None, epsilon=None):
-        """Constructs an Annealed Langevin Solver.
-        Args:
-            sde: A valid SDE class.
-            snr: A hyperparameter representing a signal-to-noise ratio.
-        """
-        super().__init__(num_steps, dt=dt, epsilon=epsilon)
-        self.sde = sde
-        self.snr = snr
-        self.prior = sde.prior
-
-    def update(self, rng, x, t):
-        grad = self.sde.score(x, t)
-        grad_norm = jnp.linalg.norm(
-            grad.reshape((grad.shape[0], -1)), axis=-1).mean()
-        grad_norm = jax.lax.pmean(grad_norm, axis_name='batch')
-        noise = random.normal(rng, x.shape)
-        noise_norm = jnp.linalg.norm(
-            noise.reshape((noise.shape[0], -1)), axis=-1).mean()
-        noise_norm = jax.lax.pmean(noise_norm, axis_name='batch')
-        alpha = jnp.exp(2 * self.sde.log_mean_coeff(t))
-        dt = (self.snr * noise_norm / grad_norm)**2 * 2 * alpha
-        x_mean = x + batch_mul(grad, dt)
-        x = x_mean + batch_mul(2 * dt, noise)
-        return x, x_mean
+  """Annealed Langevin numerical solver of an SDE.
+  Functions are designed for a mini-batch of inputs."""
+
+  def __init__(self, sde, snr=1e-2, num_steps=2, dt=None, epsilon=None):
+    """Constructs an Annealed Langevin Solver.
+    Args:
+      sde: A valid SDE class.
+      snr: A hyperparameter representing a signal-to-noise ratio.
+    """
+    super().__init__(num_steps, dt=dt, epsilon=epsilon)
+    self.sde = sde
+    self.snr = snr
+    self.prior = sde.prior
+
+  def update(self, rng, x, t):
+    grad = self.sde.score(x, t)
+    grad_norm = jnp.linalg.norm(
+      grad.reshape((grad.shape[0], -1)), axis=-1).mean()
+    grad_norm = jax.lax.pmean(grad_norm, axis_name='batch')
+    noise = random.normal(rng, x.shape)
+    noise_norm = jnp.linalg.norm(
+      noise.reshape((noise.shape[0], -1)), axis=-1).mean()
+    noise_norm = jax.lax.pmean(noise_norm, axis_name='batch')
+    alpha = jnp.exp(2 * self.sde.log_mean_coeff(t))
+    dt = (self.snr * noise_norm / grad_norm)**2 * 2 * alpha
+    x_mean = x + batch_mul(grad, dt)
+    x = x_mean + batch_mul(2 * dt, noise)
+    return x, x_mean
 
 
 class DDPM(Solver):
-    """DDPM Markov chain using Ancestral sampling."""
-    def __init__(self, score, num_steps=1000, dt=None, epsilon=None, beta_min=0.1, beta_max=20.0):
-        super().__init__(num_steps, dt, epsilon)
-        self.score = score
-        self.discrete_betas = jnp.linspace(
-            beta_min / num_steps, beta_max / num_steps, num_steps)
-        self.alphas = 1. - self.discrete_betas
-        self.alphas_cumprod = jnp.cumprod(self.alphas, axis=0)
-        self.sqrt_alphas_cumprod = jnp.sqrt(self.alphas_cumprod)
-        self.sqrt_1m_alphas_cumprod = jnp.sqrt(1. - self.alphas_cumprod)
-        self.alphas_cumprod_prev = jnp.append(1.0, self.alphas_cumprod[:-1])
-        self.sqrt_alphas_cumprod_prev = jnp.sqrt(self.alphas_cumprod_prev)
-        self.sqrt_1m_alphas_cumprod_prev = jnp.sqrt(1. - self.alphas_cumprod_prev)
-
-    def prior(self, rng, shape):
-        return random.normal(rng, shape)
-
-    def posterior(self, score, x, timestep):
-        beta = self.discrete_betas[timestep]
-        # As implemented by Song
-        # https://github.com/yang-song/score_sde/blob/0acb9e0ea3b8cccd935068cd9c657318fbc6ce4c/sampling.py#L237C5-L237C79
-        # x_mean = batch_mul(
-        #     (x + batch_mul(beta, score)), 1. / jnp.sqrt(1. - beta))  # DDPM
-        # std = jnp.sqrt(beta)
-
-        # # As implemented by DPS2022
-        # https://github.com/DPS2022/diffusion-posterior-sampling/blob/effbde7325b22ce8dc3e2c06c160c021e743a12d/guided_diffusion/gaussian_diffusion.py#L373
-        m = self.sqrt_alphas_cumprod[timestep]
-        v = self.sqrt_1m_alphas_cumprod[timestep]**2
-        alpha = self.alphas[timestep]
-        x_0 = batch_mul((x + batch_mul(v, score)), 1. / m)
-        m_prev = self.sqrt_alphas_cumprod_prev[timestep]
-        v_prev = self.sqrt_1m_alphas_cumprod_prev[timestep]**2
-        x_mean = batch_mul(jnp.sqrt(alpha) * v_prev / v, x) + batch_mul(m_prev * beta / v, x_0)
-        std = jnp.sqrt(beta * v_prev / v)
-        return x_mean, std
-
-    def update(self, rng, x, t):
-        score = self.score(x, t)
-        timestep = (t * (self.num_steps - 1) / self.t1).astype(jnp.int32)
-        x_mean, std = self.posterior(score, x, timestep)
-        z = random.normal(rng, x.shape)
-        x = x_mean + batch_mul(std, z)
-        return x, x_mean
+  """DDPM Markov chain using Ancestral sampling."""
+  def __init__(self, score, num_steps=1000, dt=None, epsilon=None, beta_min=0.1, beta_max=20.0):
+    super().__init__(num_steps, dt, epsilon)
+    self.score = score
+    self.discrete_betas = jnp.linspace(
+      beta_min / num_steps, beta_max / num_steps, num_steps)
+    self.alphas = 1. - self.discrete_betas
+    self.alphas_cumprod = jnp.cumprod(self.alphas, axis=0)
+    self.sqrt_alphas_cumprod = jnp.sqrt(self.alphas_cumprod)
+    self.sqrt_1m_alphas_cumprod = jnp.sqrt(1. - self.alphas_cumprod)
+    self.alphas_cumprod_prev = jnp.append(1.0, self.alphas_cumprod[:-1])
+    self.sqrt_alphas_cumprod_prev = jnp.sqrt(self.alphas_cumprod_prev)
+    self.sqrt_1m_alphas_cumprod_prev = jnp.sqrt(1. - self.alphas_cumprod_prev)
+
+  def get_estimate_x_0(self, shape, observation_map):
+    def estimate_x_0(x, t, timestep):
+      m = self.sqrt_alphas_cumprod[timestep]
+      v = self.sqrt_1m_alphas_cumprod[timestep]**2
+      x = x.reshape(shape)
+      x = jnp.expand_dims(x, axis=0)
+      t = jnp.expand_dims(t, axis=0)
+      s = self.score(x, t)
+      s = s.flatten()
+      x = x.flatten()
+      x_0 = (x + v * s) / m
+      return observation_map(x_0), (s, x_0)
+    return estimate_x_0
+
+  def prior(self, rng, shape):
+    return random.normal(rng, shape)
+
+  def posterior(self, score, x, timestep):
+    beta = self.discrete_betas[timestep]
+    # As implemented by Song
+    # https://github.com/yang-song/score_sde/blob/0acb9e0ea3b8cccd935068cd9c657318fbc6ce4c/sampling.py#L237C5-L237C79
+    # x_mean = batch_mul(
+    #     (x + batch_mul(beta, score)), 1. / jnp.sqrt(1. - beta))  # DDPM
+    # std = jnp.sqrt(beta)
+
+    # # As implemented by DPS2022
+    # https://github.com/DPS2022/diffusion-posterior-sampling/blob/effbde7325b22ce8dc3e2c06c160c021e743a12d/guided_diffusion/gaussian_diffusion.py#L373
+    m = self.sqrt_alphas_cumprod[timestep]
+    v = self.sqrt_1m_alphas_cumprod[timestep]**2
+    alpha = self.alphas[timestep]
+    x_0 = batch_mul((x + batch_mul(v, score)), 1. / m)
+    m_prev = self.sqrt_alphas_cumprod_prev[timestep]
+    v_prev = self.sqrt_1m_alphas_cumprod_prev[timestep]**2
+    x_mean = batch_mul(jnp.sqrt(alpha) * v_prev / v, x) + batch_mul(m_prev * beta / v, x_0)
+    std = jnp.sqrt(beta * v_prev / v)
+    return x_mean, std
+
+  def update(self, rng, x, t):
+    score = self.score(x, t)
+    timestep = (t * (self.num_steps - 1) / self.t1).astype(jnp.int32)
+    x_mean, std = self.posterior(score, x, timestep)
+    z = random.normal(rng, x.shape)
+    x = x_mean + batch_mul(std, z)
+    return x, x_mean
 
 
 class SMLD(Solver):
-    """SMLD(NCSN) Markov Chain using Ancestral sampling."""
-    def __init__(self, score, num_steps=1000, dt=None, epsilon=None, sigma_min=0.01, sigma_max=378.):
-        super().__init__(num_steps, dt, epsilon)
-        self.score = score
-        self.sigma_min = sigma_min
-        self.sigma_max = sigma_max
-        self.discrete_sigmas = jnp.exp(
-            jnp.linspace(jnp.log(self.sigma_min),
-                        jnp.log(self.sigma_max),
-                        self.num_steps))
-        self.discrete_sigmas_prev = jnp.append(0.0, self.discrete_sigmas[:-1])
-
-    def prior(self, rng, shape):
-        return random.normal(rng, shape) * self.sigma_max
-
-    def posterior(self, score, x, timestep):
-        sigma = self.discrete_sigmas[timestep]
-        sigma_prev = self.discrete_sigmas_prev[timestep]
-
-        # As implemented by Song https://github.com/yang-song/score_sde/blob/0acb9e0ea3b8cccd935068cd9c657318fbc6ce4c/sampling.py#L220
-        # x_mean = x + batch_mul(score, sigma**2 - sigma_prev**2)
-        # std = jnp.sqrt((sigma_prev**2 * (sigma**2 - sigma_prev**2)) / (sigma**2))
-
-        # From posterior in Appendix F https://arxiv.org/pdf/2011.13456.pdf
-        x_0 = x + batch_mul(sigma**2, score)
-        x_mean = batch_mul(sigma_prev**2 / sigma**2, x) + batch_mul(1 - sigma_prev**2 / sigma**2, x_0)
-        std = jnp.sqrt((sigma_prev**2 * (sigma**2 - sigma_prev**2)) / (sigma**2))
-        return x_mean, std
-
-    def update(self, rng, x, t):
-        timestep = (t * (self.num_steps - 1) / self.t1).astype(jnp.int32)
-        score = self.score(x, t)
-        x_mean, std = self.posterior(score, x, timestep)
-        z = random.normal(rng, x.shape)
-        x = x_mean + batch_mul(std, z)
-        return x, x_mean
+  """SMLD(NCSN) Markov Chain using Ancestral sampling."""
+  def __init__(self, score, num_steps=1000, dt=None, epsilon=None, sigma_min=0.01, sigma_max=378.):
+    super().__init__(num_steps, dt, epsilon)
+    self.score = score
+    self.sigma_min = sigma_min
+    self.sigma_max = sigma_max
+    self.discrete_sigmas = jnp.exp(
+        jnp.linspace(jnp.log(self.sigma_min),
+                     jnp.log(self.sigma_max),
+                     self.num_steps))
+    self.discrete_sigmas_prev = jnp.append(0.0, self.discrete_sigmas[:-1])
+
+  def get_estimate_x_0(self, shape, observation_map):
+    def estimate_x_0(x, t, timestep):
+      v = self.discrete_sigmas[timestep]**2
+      x = x.reshape(shape)
+      x = jnp.expand_dims(x, axis=0)
+      t = jnp.expand_dims(t, axis=0)
+      s = self.score(x, t)
+      x = x.flatten()
+      s = s.flatten()
+      x_0 = x + v * s
+      return observation_map(x_0), (s, x_0)
+    return estimate_x_0
+
+  def prior(self, rng, shape):
+    return random.normal(rng, shape) * self.sigma_max
+
+  def posterior(self, score, x, timestep):
+    sigma = self.discrete_sigmas[timestep]
+    sigma_prev = self.discrete_sigmas_prev[timestep]
+
+    # As implemented by Song https://github.com/yang-song/score_sde/blob/0acb9e0ea3b8cccd935068cd9c657318fbc6ce4c/sampling.py#L220
+    # x_mean = x + batch_mul(score, sigma**2 - sigma_prev**2)
+    # std = jnp.sqrt((sigma_prev**2 * (sigma**2 - sigma_prev**2)) / (sigma**2))
+
+    # From posterior in Appendix F https://arxiv.org/pdf/2011.13456.pdf
+    x_0 = x + batch_mul(sigma**2, score)
+    x_mean = batch_mul(sigma_prev**2 / sigma**2, x) + batch_mul(1 - sigma_prev**2 / sigma**2, x_0)
+    std = jnp.sqrt((sigma_prev**2 * (sigma**2 - sigma_prev**2)) / (sigma**2))
+    return x_mean, std
+
+  def update(self, rng, x, t):
+    timestep = (t * (self.num_steps - 1) / self.t1).astype(jnp.int32)
+    score = self.score(x, t)
+    x_mean, std = self.posterior(score, x, timestep)
+    z = random.normal(rng, x.shape)
+    x = x_mean + batch_mul(std, z)
+    return x, x_mean
 
 
 class DDIMVP(Solver):
-    """DDIM Markov chain. For the DDPM Markov Chain or VP SDE."""
-    def __init__(self, model, eta=0.0, num_steps=1000, dt=None, epsilon=None, beta_min=0.1, beta_max=20.0):
-        """
-        Args:
-            model: DDIM parameterizes the `\epsilon(x, t) = -1. * fwd_marginal_std(t) * score(x, t)` function.
-            eta: the hyperparameter for DDIM, a value of `eta=0.0` is deterministic 'probability ODE' solver, `eta=1.0` is DDPMVP.
-        """
-        super().__init__(num_steps, dt, epsilon)
-        self.eta = eta
-        self.model = model
-        self.beta_min = beta_min
-        self.beta_max = beta_max
-        self.discrete_betas = jnp.linspace(
-            beta_min / num_steps, beta_max / num_steps, num_steps)
-        self.alphas = 1. - self.discrete_betas
-        self.alphas_cumprod = jnp.cumprod(self.alphas, axis=0)
-        self.sqrt_alphas_cumprod = jnp.sqrt(self.alphas_cumprod)
-        self.sqrt_1m_alphas_cumprod = jnp.sqrt(1. - self.alphas_cumprod)
-        self.alphas_cumprod_prev = jnp.append(1.0, self.alphas_cumprod[:-1])
-        self.sqrt_alphas_cumprod_prev = jnp.sqrt(self.alphas_cumprod_prev)
-        self.sqrt_1m_alphas_cumprod_prev = jnp.sqrt(1. - self.alphas_cumprod_prev)
-
-    def prior(self, rng, shape):
-        return random.normal(rng, shape)
-
-    def posterior(self, x, t):
-        # # As implemented by DPS2022
-        # https://github.com/DPS2022/diffusion-posterior-sampling/blob/effbde7325b22ce8dc3e2c06c160c021e743a12d/guided_diffusion/gaussian_diffusion.py#L373
-        # and as written in https://arxiv.org/pdf/2010.02502.pdf
-        epsilon = self.model(x, t)
-        timestep = (t * (self.num_steps - 1) / self.t1).astype(jnp.int32)
-        m = self.sqrt_alphas_cumprod[timestep]
-        sqrt_1m_alpha = self.sqrt_1m_alphas_cumprod[timestep]
-        v = sqrt_1m_alpha**2
-        alpha = m**2
-        m_prev = self.sqrt_alphas_cumprod_prev[timestep]
-        v_prev = self.sqrt_1m_alphas_cumprod_prev[timestep]**2
-        alpha_prev = m_prev**2
-        x_0 = batch_mul((x - batch_mul(sqrt_1m_alpha, epsilon)), 1. / m)
-        coeff1 = self.eta * jnp.sqrt((v_prev / v) * (1 - alpha / alpha_prev))
-        coeff2 = jnp.sqrt(v_prev - coeff1**2)
-        x_mean = batch_mul(m_prev, x_0) + batch_mul(coeff2, epsilon)
-        std = coeff1
-        return x_mean, std
-
-    def update(self, rng, x, t):
-        x_mean, std = self.posterior(x, t)
-        z = random.normal(rng, x.shape)
-        x = x_mean + batch_mul(std, z)
-        return x, x_mean
+  """DDIM Markov chain. For the DDPM Markov Chain or VP SDE."""
+  def __init__(self, model, eta=0.0, num_steps=1000, dt=None, epsilon=None, beta_min=0.1, beta_max=20.0):
+    """
+    Args:
+        model: DDIM parameterizes the `\epsilon(x, t) = -1. * fwd_marginal_std(t) * score(x, t)` function.
+        eta: the hyperparameter for DDIM, a value of `eta=0.0` is deterministic 'probability ODE' solver, `eta=1.0` is DDPMVP.
+    """
+    super().__init__(num_steps, dt, epsilon)
+    self.eta = eta
+    self.model = model
+    self.beta_min = beta_min
+    self.beta_max = beta_max
+    self.discrete_betas = jnp.linspace(
+      beta_min / num_steps, beta_max / num_steps, num_steps)
+    self.alphas = 1. - self.discrete_betas
+    self.alphas_cumprod = jnp.cumprod(self.alphas, axis=0)
+    self.sqrt_alphas_cumprod = jnp.sqrt(self.alphas_cumprod)
+    self.sqrt_1m_alphas_cumprod = jnp.sqrt(1. - self.alphas_cumprod)
+    self.alphas_cumprod_prev = jnp.append(1.0, self.alphas_cumprod[:-1])
+    self.sqrt_alphas_cumprod_prev = jnp.sqrt(self.alphas_cumprod_prev)
+    self.sqrt_1m_alphas_cumprod_prev = jnp.sqrt(1. - self.alphas_cumprod_prev)
+
+  def get_estimate_x_0(self, shape, observation_map):
+    def estimate_x_0(x, t, timestep):
+      m = self.sqrt_alphas_cumprod[timestep]
+      sqrt_1m_alpha = self.sqrt_1m_alphas_cumprod[timestep]
+      x = x.reshape(shape)
+      x = jnp.expand_dims(x, axis=0)
+      t = jnp.expand_dims(t, axis=0)
+      epsilon = self.model(x, t)
+      epsilon = epsilon.flatten()
+      x = x.flatten()
+      x_0 = (x - sqrt_1m_alpha * epsilon) / m
+      return observation_map(x_0), (epsilon, x_0)
+    return estimate_x_0
+
+  def prior(self, rng, shape):
+    return random.normal(rng, shape)
+
+  def posterior(self, x, t):
+    # # As implemented by DPS2022
+    # https://github.com/DPS2022/diffusion-posterior-sampling/blob/effbde7325b22ce8dc3e2c06c160c021e743a12d/guided_diffusion/gaussian_diffusion.py#L373
+    # and as written in https://arxiv.org/pdf/2010.02502.pdf
+    epsilon = self.model(x, t)
+    timestep = (t * (self.num_steps - 1) / self.t1).astype(jnp.int32)
+    m = self.sqrt_alphas_cumprod[timestep]
+    sqrt_1m_alpha = self.sqrt_1m_alphas_cumprod[timestep]
+    v = sqrt_1m_alpha**2
+    alpha = m**2
+    m_prev = self.sqrt_alphas_cumprod_prev[timestep]
+    v_prev = self.sqrt_1m_alphas_cumprod_prev[timestep]**2
+    alpha_prev = m_prev**2
+    x_0 = batch_mul((x - batch_mul(sqrt_1m_alpha, epsilon)), 1. / m)
+    coeff1 = self.eta * jnp.sqrt((v_prev / v) * (1 - alpha / alpha_prev))
+    coeff2 = jnp.sqrt(v_prev - coeff1**2)
+    x_mean = batch_mul(m_prev, x_0) + batch_mul(coeff2, epsilon)
+    std = coeff1
+    return x_mean, std
+
+  def update(self, rng, x, t):
+    x_mean, std = self.posterior(x, t)
+    z = random.normal(rng, x.shape)
+    x = x_mean + batch_mul(std, z)
+    return x, x_mean
 
 
 class DDIMVE(Solver):
-    """DDIM Markov chain. For the SMLD Markov Chain or VE SDE.
+  """DDIM Markov chain. For the SMLD Markov Chain or VE SDE.
+  """
+  def __init__(self, model, eta=0.0, num_steps=1000, dt=None, epsilon=None, sigma_min=0.01, sigma_max=378.):
+    """
+    Args:
+        model: DDIM parameterizes the `\epsilon(x, t) = -1. * fwd_marginal_std(t) * score(x, t)` function.
+        eta: the hyperparameter for DDIM, a value of `eta=0.0` is deterministic 'probability ODE' solver, `eta=1.0` is DDPMVE.
     """
-    def __init__(self, model, eta=0.0, num_steps=1000, dt=None, epsilon=None, sigma_min=0.01, sigma_max=378.):
-        """
-        Args:
-            model: DDIM parameterizes the `\epsilon(x, t) = -1. * fwd_marginal_std(t) * score(x, t)` function.
-            eta: the hyperparameter for DDIM, a value of `eta=0.0` is deterministic 'probability ODE' solver, `eta=1.0` is DDPMVE.
-        """
-        super().__init__(num_steps, dt, epsilon)
-        self.eta = eta
-        self.model = model
-        self.sigma_min = sigma_min
-        self.sigma_max = sigma_max
-        self.discrete_sigmas = jnp.exp(
-            jnp.linspace(jnp.log(self.sigma_min),
-                        jnp.log(self.sigma_max),
-                        self.num_steps))
-        self.discrete_sigmas_prev = jnp.append(0.0, self.discrete_sigmas[:-1])
-
-    def prior(self, rng, shape):
-        return random.normal(rng, shape) * self.sigma_max
-
-    def posterior(self, x, t):
-        timestep = (t * (self.num_steps - 1) / self.t1).astype(jnp.int32)
-        epsilon = self.model(x, t)
-        sigma = self.discrete_sigmas[timestep]
-        sigma_prev = self.discrete_sigmas_prev[timestep]
-
-        # Eq.(18) Appendix A.4 https://openreview.net/pdf/210093330709030207aa90dbfe2a1f525ac5fb7d.pdf
-        coeff1 = self.eta * jnp.sqrt((sigma_prev**2 * (sigma**2 - sigma_prev**2)) / (sigma**2))
-        coeff2 = jnp.sqrt(sigma_prev**2  - coeff1**2)
-        x_0 = x - batch_mul(sigma, epsilon)  # TODO: check
-        x_mean = x_0 + batch_mul(coeff2, epsilon)
-        std = coeff1
-        return x_mean, std
-
-    def update(self, rng, x, t):
-        x_mean, std = self.posterior(x, t)
-        z = random.normal(rng, x.shape)
-        x = x_mean + batch_mul(std, z)
-        return x, x_mean
+    super().__init__(num_steps, dt, epsilon)
+    self.eta = eta
+    self.model = model
+    self.sigma_min = sigma_min
+    self.sigma_max = sigma_max
+    self.discrete_sigmas = jnp.exp(
+        jnp.linspace(jnp.log(self.sigma_min),
+                    jnp.log(self.sigma_max),
+                    self.num_steps))
+    self.discrete_sigmas_prev = jnp.append(0.0, self.discrete_sigmas[:-1])
+
+  def get_estimate_x_0(self, shape, observation_map):
+    def estimate_x_0(x, t, timestep):
+      std = self.discrete_sigmas[timestep]
+      x = x.reshape(shape)
+      x = jnp.expand_dims(x, axis=0)
+      t = jnp.expand_dims(t, axis=0)
+      epsilon = self.model(x, t)
+      epsilon = epsilon.flatten()
+      x = x.flatten()
+      x_0 = x - std * epsilon
+      return observation_map(x_0), (epsilon, x_0)
+    return estimate_x_0
+
+  def prior(self, rng, shape):
+    return random.normal(rng, shape) * self.sigma_max
+
+  def posterior(self, x, t):
+    timestep = (t * (self.num_steps - 1) / self.t1).astype(jnp.int32)
+    epsilon = self.model(x, t)
+    sigma = self.discrete_sigmas[timestep]
+    sigma_prev = self.discrete_sigmas_prev[timestep]
+    coeff1 = self.eta * jnp.sqrt((sigma_prev**2 * (sigma**2 - sigma_prev**2)) / (sigma**2))
+    coeff2 = jnp.sqrt(sigma_prev**2  - coeff1**2)
+
+    # Eq.(18) Appendix A.4 https://openreview.net/pdf/210093330709030207aa90dbfe2a1f525ac5fb7d.pdf
+    x_0 = x - batch_mul(sigma, epsilon)
+    x_mean = x_0 + batch_mul(coeff2, epsilon)
+
+    # Eq.(19) Appendix A.4 https://openreview.net/pdf/210093330709030207aa90dbfe2a1f525ac5fb7d.pdf
+    # score = - batch_mul(1. / sigma, epsilon)
+    # x_mean = x + batch_mul(sigma * (sigma - coeff2), score)
+
+    std = coeff1
+    return x_mean, std
+
+  def update(self, rng, x, t):
+    x_mean, std = self.posterior(x, t)
+    z = random.normal(rng, x.shape)
+    x = x_mean + batch_mul(std, z)
+    return x, x_mean
diff --git a/diffusionjax/utils.py b/diffusionjax/utils.py
index 1533ccb..356bcfc 100644
--- a/diffusionjax/utils.py
+++ b/diffusionjax/utils.py
@@ -1,395 +1,183 @@
 """Utility functions, including all functions related to
-loss computation, optimization, sampling and inverse problems.
+loss computation, optimization and sampling.
 """
 import jax.numpy as jnp
 from jax.lax import scan
-from jax import vmap, jit, value_and_grad
+from jax import vmap
 import jax.random as random
 from functools import partial
-import optax
-import numpy as np
 
 
 def batch_mul(a, b):
-    return vmap(lambda a, b: a * b)(a, b)
+  return vmap(lambda a, b: a * b)(a, b)
 
 
 def errors(ts, sde, score, rng, data, likelihood_weighting=True):
-    """
-    Args:
-        ts: JAX array of times.
-        sde: Instantiation of a valid SDE class.
-        score: A function taking in (x, t) and returning the score.
-        rng: Random number generator from JAX.
-        data: A batch of samples from the training data, representing samples from \mu_text{data}, shape (J, N).
-        likelihood_weighting: Boolean variable, set to `True` if likelihood weighting, as described in Song et al. 2020 (https://arxiv.org/abs/2011.13456), is applied.
-    Returns:
-        A random (MC) approximation to the (likelihood weighted) score errors.
-    """
-    mean, std = sde.marginal_prob(data, ts)
-    rng, step_rng = random.split(rng)
-    noise = random.normal(step_rng, data.shape)
-    x_t = mean + batch_mul(std, noise)
-    if not likelihood_weighting:
-        return noise + batch_mul(score(x_t, ts), std)
-    else:
-        return batch_mul(noise, 1. / std) + score(x_t, ts)
+  """
+  Args:
+    ts: JAX array of times.
+    sde: Instantiation of a valid SDE class.
+    score: A function taking in (x, t) and returning the score.
+    rng: Random number generator from JAX.
+    data: A batch of samples from the training data, representing samples from \mu_text{data}, shape (J, N).
+    likelihood_weighting: Boolean variable, set to `True` if likelihood weighting, as described in Song et al. 2020 (https://arxiv.org/abs/2011.13456), is applied.
+  Returns:
+    A random (MC) approximation to the (likelihood weighted) score errors.
+  """
+  mean, std = sde.marginal_prob(data, ts)
+  rng, step_rng = random.split(rng)
+  noise = random.normal(step_rng, data.shape)
+  x_t = mean + batch_mul(std, noise)
+  if not likelihood_weighting:
+    return noise + batch_mul(score(x_t, ts), std)
+  else:
+    return batch_mul(noise, 1. / std) + score(x_t, ts)
 
 
 def get_loss(sde, solver, model, score_scaling=True, likelihood_weighting=True, reduce_mean=True, pointwise_t=False):
-    """Create a loss function for score matching training.
-    Args:
-        sde: Instantiation of a valid SDE class.
-        solver: Instantiation of a valid Solver class.
-        model: A valid flax neural network `:class:flax.linen.Module` class.
-        score_scaling: Boolean variable, set to `True` if learning a score scaled by the marginal standard deviation.
-        likelihood_weighting: Boolean variable, set to `True` if likelihood weighting, as described in Song et al. 2020 (https://arxiv.org/abs/2011.13456), is applied.
-        reduce_mean: Boolean variable, set to `True` if taking the mean of the errors in the loss, set to `False` if taking the sum.
-        pointwise_t: Boolean variable, set to `True` if returning a function that can evaluate the loss pointwise over time. Set to `False` if returns an expectation of the loss over time.
-
-    Returns:
-        A loss function that can be used for score matching training.
-    """
-    reduce_op = jnp.mean if reduce_mean else lambda *args, **kwargs: 0.5 * jnp.sum(*args, **kwargs)
-    if pointwise_t:
-        def loss(t, params, rng, data):
-            n_batch = data.shape[0]
-            ts = jnp.ones((n_batch,)) * t
-            score = get_score(sde, model, params, score_scaling)
-            e = errors(ts, sde, score, rng, data, likelihood_weighting)
-            losses = e**2
-            losses = reduce_op(losses.reshape((losses.shape[0], -1)), axis=-1)
-            if likelihood_weighting:
-                g2 = sde.sde(jnp.zeros_like(data), ts)[1]**2
-                losses = losses * g2
-            return jnp.mean(losses)
-    else:
-        def loss(params, rng, data):
-            rng, step_rng = random.split(rng)
-            n_batch = data.shape[0]
-            ts = random.uniform(step_rng, (data.shape[0],), minval=solver.ts[0], maxval=solver.t1)
-            score = get_score(sde, model, params, score_scaling)
-            e = errors(ts, sde, score, rng, data, likelihood_weighting)
-            losses = e**2
-            losses = reduce_op(losses.reshape((losses.shape[0], -1)), axis=-1)
-            if likelihood_weighting:
-                g2 = sde.sde(jnp.zeros_like(data), ts)[1]**2
-                losses = losses * g2
-            return jnp.mean(losses)
-    return loss
-
-
-def get_step_fn(loss, optimizer, train, pmap):
-    """Create a one-step training/evaluation function.
-
-    Args:
-        loss: A loss function.
-        optimizer: An optimization function.
-        train: `True` for training and `False` for evaluation.
-        pmap: `True` for pmap across jax devices, `False` for single device.
-
-    Returns:
-        A one-step function for training or evaluation.
-    """
-    @jit
-    def step_fn(carry, batch):
-        (rng, params, opt_state) = carry
-        rng, step_rng = random.split(rng)
-        grad_fn = value_and_grad(loss)
-        if train:
-            loss_val, grads = grad_fn(params, step_rng, batch)
-            if pmap:
-                loss_val = jax.lax.pmean(loss_val, axis_name='batch')
-                grads = jax.lax.pmean(grads, axis_name='batch')
-            updates, opt_state = optimizer.update(grads, opt_state)
-            params = optax.apply_updates(params, updates)
-        else:
-            loss_val = loss(params, step_rng, batch)
-            if pmap: loss_val = jax.lax.pmean(loss_val, axis_name='batch')
-        return (rng, params, opt_state), loss_val
-    return step_fn
-
-
-#Initialize the optimizer
-optimizer = optax.adam(1e-3)
-
-
-@partial(jit, static_argnums=[4])
-def update_step(params, rng, batch, opt_state, loss):
-    """
-    Takes the gradient of the loss function and updates the model weights (params) using it.
-    Args:
-        params: the current weights of the model
-        rng: random number generator from jax
-        batch: a batch of samples from the training data, representing samples from \mu_text{data}, shape (J, N)
-        opt_state: the internal state of the optimizer
-        loss: A loss function that can be used for score matching training.
-    Returns:
-        The value of the loss function (for metrics), the new params and the new optimizer states function (for metrics),
-        the new params and the new optimizer state.
-    """
-    val, grads = value_and_grad(loss)(params, rng, batch)
-    updates, opt_state = optimizer.update(grads, opt_state)
-    params = optax.apply_updates(params, updates)
-    return val, params, opt_state
-
-
-def retrain_nn(
-        update_step, num_epochs, step_rng, samples, params,
-        opt_state, loss, batch_size=5):
-    train_size = samples.shape[0]
-    batch_size = min(train_size, batch_size)
-    steps_per_epoch = train_size // batch_size
-    mean_losses = jnp.zeros((num_epochs, 1))
-    for i in range(num_epochs):
-        rng, step_rng = random.split(step_rng)
-        perms = random.permutation(step_rng, train_size)
-        perms = perms[:steps_per_epoch * batch_size]  # skip incomplete batch
-        perms = perms.reshape((steps_per_epoch, batch_size))
-        losses = jnp.zeros((jnp.shape(perms)[0], 1))
-        for j, perm in enumerate(perms):
-            batch = samples[perm, :]
-            rng, step_rng = random.split(rng)
-            loss_eval, params, opt_state = update_step(params, step_rng, batch, opt_state, loss)
-            losses = losses.at[j].set(loss_eval)
-        mean_loss = jnp.mean(losses, axis=0)
-        mean_losses = mean_losses.at[i].set(mean_loss)
-        if i % 10 == 0:
-            print("Epoch {:d}, Loss {:.2f} ".format(i, mean_loss[0]))
-    return params, opt_state, mean_losses
+  """Create a loss function for score matching training.
+  Args:
+    sde: Instantiation of a valid SDE class.
+    solver: Instantiation of a valid Solver class.
+    model: A valid flax neural network `:class:flax.linen.Module` class.
+    score_scaling: Boolean variable, set to `True` if learning a score scaled by the marginal standard deviation.
+    likelihood_weighting: Boolean variable, set to `True` if likelihood weighting, as described in Song et al. 2020 (https://arxiv.org/abs/2011.13456), is applied.
+    reduce_mean: Boolean variable, set to `True` if taking the mean of the errors in the loss, set to `False` if taking the sum.
+    pointwise_t: Boolean variable, set to `True` if returning a function that can evaluate the loss pointwise over time. Set to `False` if returns an expectation of the loss over time.
+
+  Returns:
+    A loss function that can be used for score matching training.
+  """
+  reduce_op = jnp.mean if reduce_mean else lambda *args, **kwargs: 0.5 * jnp.sum(*args, **kwargs)
+  if pointwise_t:
+    def loss(t, params, rng, data):
+      n_batch = data.shape[0]
+      ts = jnp.ones((n_batch,)) * t
+      score = get_score(sde, model, params, score_scaling)
+      e = errors(ts, sde, score, rng, data, likelihood_weighting)
+      losses = e**2
+      losses = reduce_op(losses.reshape((losses.shape[0], -1)), axis=-1)
+      if likelihood_weighting:
+        g2 = sde.sde(jnp.zeros_like(data), ts)[1]**2
+        losses = losses * g2
+      return jnp.mean(losses)
+  else:
+    def loss(params, rng, data):
+      rng, step_rng = random.split(rng)
+      ts = random.uniform(step_rng, (data.shape[0],), minval=solver.ts[0], maxval=solver.t1)
+      score = get_score(sde, model, params, score_scaling)
+      e = errors(ts, sde, score, rng, data, likelihood_weighting)
+      losses = e**2
+      losses = reduce_op(losses.reshape((losses.shape[0], -1)), axis=-1)
+      if likelihood_weighting:
+        g2 = sde.sde(jnp.zeros_like(data), ts)[1]**2
+        losses = losses * g2
+      return jnp.mean(losses)
+  return loss
 
 
 def get_score(sde, model, params, score_scaling):
-    if score_scaling is True:
-        return lambda x, t: -batch_mul(model.apply(params, x, t), 1. / jnp.sqrt(sde.variance(t)))
-    else:
-        return lambda x, t: -model.apply(params, x, t)
+  if score_scaling is True:
+    return lambda x, t: -batch_mul(model.apply(params, x, t), 1. / jnp.sqrt(sde.variance(t)))
+  else:
+    return lambda x, t: -model.apply(params, x, t)
 
 
 def get_epsilon(sde, model, params, score_scaling):
-    if score_scaling is True:
-        return lambda x, t: model.apply(params, x, t)
-    else:
-        return lambda x, t: batch_mul(jnp.sqrt(sde.variance(t)), model.apply(params, x, t))
+  if score_scaling is True:
+    return lambda x, t: model.apply(params, x, t)
+  else:
+    return lambda x, t: batch_mul(jnp.sqrt(sde.variance(t)), model.apply(params, x, t))
 
 
 def shared_update(rng, x, t, solver, probability_flow=None):
-    """A wrapper that configures and returns the update function of the solvers.
+  """A wrapper that configures and returns the update function of the solvers.
 
-    :probablity_flow: Placeholder for probability flow ODE (TODO).
-    """
-    return solver.update(rng, x, t)
+  :probablity_flow: Placeholder for probability flow ODE (TODO).
+  """
+  return solver.update(rng, x, t)
 
 
 def get_sampler(shape, outer_solver, inner_solver=None, denoise=True, stack_samples=False, inverse_scaler=None):
-    """Get a sampler from (possibly interleaved) numerical solver(s).
-
-    Args:
-        shape: Shape of array, x. (num_samples,) + obj_shape, where x_shape is the shape
-            of the object being sampled from, for example, an image may have
-            obj_shape==(H, W, C), and so shape==(N, H, W, C) where N is the number of samples.
-        outer_solver: A valid numerical solver class that will act on an outer loop.
-        inner_solver: '' that will act on an inner loop.
-        denoise: Boolean variable that if `True` applies one-step denoising to final samples.
-        stack_samples: Boolean variable that if `True` return all of the sample path or
-            just returns the last sample.
-        inverse_scaler: The inverse data normalizer function.
-    Returns:
-        A sampler.
+  """Get a sampler from (possibly interleaved) numerical solver(s).
+
+  Args:
+    shape: Shape of array, x. (num_samples,) + obj_shape, where x_shape is the shape
+      of the object being sampled from, for example, an image may have
+      obj_shape==(H, W, C), and so shape==(N, H, W, C) where N is the number of samples.
+    outer_solver: A valid numerical solver class that will act on an outer loop.
+    inner_solver: '' that will act on an inner loop.
+    denoise: Boolean variable that if `True` applies one-step denoising to final samples.
+    stack_samples: Boolean variable that if `True` return all of the sample path or
+      just returns the last sample.
+    inverse_scaler: The inverse data normalizer function.
+  Returns:
+    A sampler.
+  """
+  if inverse_scaler is None: inverse_scaler = lambda x: x
+
+  def sampler(rng, x_0=None):
     """
-    if inverse_scaler is None: inverse_scaler = lambda x: x
-
-    def sampler(rng, x_0=None):
-        """
-        Args:
-            rng: A JAX random state.
-            x_0: Initial condition. If `None`, then samples an initial condition from the
-                sde's initial condition prior. Note that this initial condition represents
-                `x_T \sim \text{Normal}(O, I)` in reverse-time diffusion.
-        Returns:
-            Samples and the number of score function (model) evaluations.
-        """
-        outer_update = partial(shared_update,
-                                         solver=outer_solver)
-        outer_ts = outer_solver.ts
-
-        if inner_solver:
-            inner_update = partial(shared_update,
-                                            solver=inner_solver)
-            inner_ts = inner_solver.ts
-            num_function_evaluations = jnp.size(outer_ts) * (jnp.size(inner_ts) + 1)
-
-            def inner_step(carry, t):
-                rng, x, x_mean, vec_t = carry
-                rng, step_rng = random.split(rng)
-                x, x_mean = inner_update(step_rng, x, vec_t)
-                return (rng, x, x_mean, vec_t), ()
-
-            def outer_step(carry, t):
-                rng, x, x_mean = carry
-                vec_t = jnp.full(shape[0], t)
-                rng, step_rng = random.split(rng)
-                x, x_mean = outer_update(step_rng, x, vec_t)
-                (rng, x, x_mean, vec_t), _ = scan(inner_step, (step_rng, x, x_mean, vec_t), inner_ts)
-                if not stack_samples:
-                    return (rng, x, x_mean), ()
-                else:
-                    if denoise:
-                        return (rng, x, x_mean), x_mean
-                    else:
-                        return (rng, x, x_mean), x
-        else:
-            num_function_evaluations = jnp.size(outer_ts)
-            def outer_step(carry, t):
-                rng, x, x_mean = carry
-                vec_t = jnp.full(shape[0], t)
-                rng, step_rng = random.split(rng)
-                x, x_mean = outer_update(step_rng, x, vec_t)
-                if not stack_samples:
-                    return (rng, x, x_mean), ()
-                else:
-                    return ((rng, x, x_mean), x_mean) if denoise else ((rng, x, x_mean), x)
-
-        rng, step_rng = random.split(rng)
-        if x_0 is None:
-            x = outer_solver.prior(step_rng, shape)
-        else:
-            assert(x_0.shape==shape)
-            x = x_0
-        if not stack_samples:
-            (_, x, x_mean), _ = scan(outer_step, (rng, x, x), outer_ts, reverse=True)
-
-            return inverse_scaler(x_mean if denoise else x), num_function_evaluations
-        else:
-            (_, _, _), xs = scan(outer_step, (rng, x, x), outer_ts, reverse=True)
-            return inverse_scaler(xs), num_function_evaluations
-    # return jax.pmap(sampler, in_axes=(0), axis_name='batch')
-    return sampler
-
-
-def merge_data_with_mask(x_space, data, mask, coeff=1.):
-    return data * mask * coeff + x_space * (1. - mask * coeff)
-
-
-def get_projection_sampler(solver, inverse_scaler=None, denoise=True, stack_samples=False):
-    """Create an image inpainting function that uses sampler, that returns only the final sample.
-
     Args:
-        inverse_scaler: The inverse data normalizer.
-        denoise: Boolean variable that if `True` applies one-step denoising to final samples.
-        stack_samples: Boolean variable that if `True` returns all samples on path(s).
+      rng: A JAX random state.
+      x_0: Initial condition. If `None`, then samples an initial condition from the
+          sde's initial condition prior. Note that this initial condition represents
+          `x_T \sim \text{Normal}(O, I)` in reverse-time diffusion.
     Returns:
-        A projection sampler function.
+        Samples and the number of score function (model) evaluations.
     """
-    if inverse_scaler is None:
-        def inverse_scaler(x):
-            return x
-    vmap_inverse_scaler = vmap(inverse_scaler)
-
-    def update(rng, data, mask, x, vec_t, coeff):
-        data_mean, std = solver.sde.marginal_prob(data, vec_t)
-        z = random.normal(rng, x.shape)
-        z_data = data_mean + batch_mul(std, z)
-        x = merge_data_with_mask(x, z_data, mask, coeff)
-
-        rng, step_rng = random.split(rng)
-        x, x_mean = solver.update(step_rng, x, vec_t)
-        return x, x_mean
-
-    ts = solver.ts
-    num_function_evaluations = jnp.size(ts)
-
-    def projection_sampler(rng, data, mask, coeff):
-        """Sampler for image inpainting.
-
-        Args:
-            rng: A JAX random state.
-            data: A JAX array that represents a mini-batch of images to inpaint.
-            mask: A {0, 1} array with the same shape of `data`. Value `1` marks known pixels,
-                and value `0` marks pixels that require inpainting.
-
-        Returns:
-            Inpainted (complete) images.
-        """
-        # Initial sample
+    outer_update = partial(shared_update,
+                           solver=outer_solver)
+    outer_ts = outer_solver.ts
+
+    if inner_solver:
+        inner_update = partial(shared_update,
+                               solver=inner_solver)
+        inner_ts = inner_solver.ts
+        num_function_evaluations = jnp.size(outer_ts) * (jnp.size(inner_ts) + 1)
+
+        def inner_step(carry, t):
+          rng, x, x_mean, vec_t = carry
+          rng, step_rng = random.split(rng)
+          x, x_mean = inner_update(step_rng, x, vec_t)
+          return (rng, x, x_mean, vec_t), ()
+
+        def outer_step(carry, t):
+          rng, x, x_mean = carry
+          vec_t = jnp.full(shape[0], t)
+          rng, step_rng = random.split(rng)
+          x, x_mean = outer_update(step_rng, x, vec_t)
+          (rng, x, x_mean, vec_t), _ = scan(inner_step, (step_rng, x, x_mean, vec_t), inner_ts)
+          if not stack_samples:
+            return (rng, x, x_mean), ()
+          else:
+            if denoise:
+              return (rng, x, x_mean), x_mean
+            else:
+              return (rng, x, x_mean), x
+    else:
+      num_function_evaluations = jnp.size(outer_ts)
+      def outer_step(carry, t):
+        rng, x, x_mean = carry
+        vec_t = jnp.full(shape[0], t)
         rng, step_rng = random.split(rng)
-        shape = data.shape
-        x = random.normal(step_rng, shape)
-
-        def f(carry, t):
-            rng, x, x_mean = carry
-            vec_t = jnp.full(shape[0], t)
-            rng, step_rng = random.split(rng)
-            x, x_mean = update(step_rng, data, mask, x, vec_t, coeff)
-            return ((rng, x, x_mean), x) if stack_samples else ((rng, x, x_mean), ())
-
+        x, x_mean = outer_update(step_rng, x, vec_t)
         if not stack_samples:
-            (_, x, _), _ = scan(f, (rng, x, x), ts, reverse=True)
-            return inverse_scaler(x), num_function_evaluations
+          return (rng, x, x_mean), ()
         else:
-            (_, _, _), xs = scan(f, (rng, x, x), ts, reverse=True)
-            return vmap_inverse_scaler(xs), num_function_evaluations
+          return ((rng, x, x_mean), x_mean) if denoise else ((rng, x, x_mean), x)
 
-    # return jax.pmap(projection_sampler, in_axes=(0, None, None, None), axis_name='batch')
-    return projection_sampler
-
-
-def get_inpainter(solver, inverse_scaler=None,
-                denoise=True, stack_samples=False):
-    """Create an image inpainting function that uses sampler, that returns only the final sample.
-
-    Args:
-        inverse_scaler: The inverse data normalizer.
-        denoise: Boolean variable that if `True` applies one-step denoising to final samples.
-    Returns:
-        An inpainting function.
-    """
-    if inverse_scaler is None:
-        def inverse_scaler(x):
-            return x
-    vmap_inverse_scaler = vmap(inverse_scaler)
-
-    def update(rng, data, mask, x, vec_t):
-        x, x_mean = solver.update(rng, x, vec_t)
-        masked_data_mean, std = solver.sde.marginal_prob(data, vec_t)
-        rng, step_rng = random.split(rng)
-        masked_data = masked_data_mean + batch_mul(random.normal(rng, x.shape), std)
-        x = x * (1. - mask) + masked_data * mask
-        x_mean = x * (1. - mask) + masked_data_mean * mask
-        return x, x_mean
-
-    ts = solver.ts
-    num_function_evaluations = jnp.size(ts)
-
-    def inpainter(rng, data, mask):
-        """Sampler for image inpainting.
-
-        Args:
-            rng: A JAX random state.
-            data: A JAX array that represents a mini-batch of images to inpaint.
-            mask: A {0, 1} array with the same shape of `data`. Value `1` marks known pixels,
-                and value `0` marks pixels that require inpainting.
-
-        Returns:
-            Inpainted (complete) images.
-        """
-        # Initial sample
-        rng, step_rng = random.split(rng)
-        shape = data.shape
-        x = data * mask + random.normal(step_rng, shape) * (1. - mask)
-
-        def f(carry, t):
-            rng, x, x_mean = carry
-            vec_t = jnp.full(shape[0], t)
-            rng, step_rng = random.split(rng)
-            x, x_mean = update(rng, data, mask, x, vec_t)
-            return ((rng, x, x_mean), x) if stack_samples else ((rng, x, x_mean), ())
-
-        if not stack_samples:
-            (_, x, _), _ = scan(f, (rng, x, x), ts, reverse=True)
-            return inverse_scaler(x), num_function_evaluations
-        else:
-            (_, _, _), xs = scan(f, (rng, x, x), ts, reverse=True)
-            return vmap_inverse_scaler(xs), num_function_evaluations
-    # return jax.pmap(inpainter, in_axes=(0, None, None), axis_name='batch')
-    return inpainter
+    rng, step_rng = random.split(rng)
+    if x_0 is None:
+      x = outer_solver.prior(step_rng, shape)
+    else:
+      assert(x_0.shape==shape)
+      x = x_0
+    if not stack_samples:
+      (_, x, x_mean), _ = scan(outer_step, (rng, x, x), outer_ts, reverse=True)
+      return inverse_scaler(x_mean if denoise else x), num_function_evaluations
+    else:
+      (_, _, _), xs = scan(outer_step, (rng, x, x), outer_ts, reverse=True)
+      return inverse_scaler(xs), num_function_evaluations
+  # return jax.pmap(sampler, in_axes=(0), axis_name='batch')
+  return sampler
diff --git a/examples/example.py b/examples/example.py
index f6e2b1c..ea7c94b 100644
--- a/examples/example.py
+++ b/examples/example.py
@@ -4,213 +4,209 @@
 A tutorial on the theoretical and implementation aspects of score-based generative models, also called diffusion models.
 """
 # Uncomment to enable double precision
-from jax.config import config as jax_config
-jax_config.update("jax_enable_x64", True)
-
+# from jax.config import config as jax_config
+# jax_config.update("jax_enable_x64", True)
 import jax
 from jax import jit, vmap, grad
 import jax.random as random
 import jax.numpy as jnp
 from jax.scipy.special import logsumexp
-
 from diffusionjax.run_lib import get_model, train, get_solver
-from diffusionjax.utils import get_score, get_sampler, get_inpainter
-from diffusionjax.plot import (plot_samples, plot_score,
-                               plot_heatmap, plot_scatter)
+from diffusionjax.utils import get_score, get_sampler
+from diffusionjax.inverse_problems import get_inpainter
+from diffusionjax.plot import plot_samples, plot_score, plot_heatmap
 import diffusionjax.sde as sde_lib
-
 from absl import app, flags
 from ml_collections.config_flags import config_flags
 from flax import serialization
 import time
-from torch.utils.data import Dataset
-import matplotlib.pyplot as plt
 import os
 
+# Dependencies:
+# This example requires optax, https://optax.readthedocs.io/en/latest/
+# This example requires orbax, https://orbax.readthedocs.io/en/latest/
+# This example requires torch[cpu], https://pytorch.org/get-started/locally/
+from torch.utils.data import Dataset
+
 
 FLAGS = flags.FLAGS
 config_flags.DEFINE_config_file(
-    "config", './configs/example.py', "Training configuration.",
-    lock_config=True)
+  "config", './configs/example.py', "Training configuration.",
+  lock_config=True)
 flags.DEFINE_string("workdir", './examples/', "Work directory.")
 flags.mark_flags_as_required(["workdir", "config"])
 
 
 class CircleDataset(Dataset):
-    """Dataset containing samples from the circle."""
-    def __init__(self, num_samples):
-        self.train_data = self.sample_circle(num_samples)
-
-    def __len__(self):
-        return self.train_data.shape[0]
+  """Dataset containing samples from the circle."""
+  def __init__(self, num_samples):
+    self.train_data = self.sample_circle(num_samples)
 
-    def __getitem__(self, idx):
-        return self.train_data[idx]
+  def __len__(self):
+    return self.train_data.shape[0]
 
-    def sample_circle(self, num_samples):
-        """Samples from the unit circle, angles split.
+  def __getitem__(self, idx):
+    return self.train_data[idx]
 
-        Args:
-            num_samples: The number of samples.
+  def sample_circle(self, num_samples):
+    """Samples from the unit circle, angles split.
 
-        Returns:
-            An (num_samples, 2) array of samples.
+    Args:
+      num_samples: The number of samples.
 
-        num_samples: Number of samples
-        Returns a (num_samples, 2) array of samples
-        """
-        alphas = jnp.linspace(0, 2 * jnp.pi * (1 - 1/num_samples), num_samples)
-        xs = jnp.cos(alphas)
-        ys = jnp.sin(alphas)
-        samples = jnp.stack([xs, ys], axis=1)
-        return samples
+    Returns:
+      An (num_samples, 2) array of samples.
+    """
+    alphas = jnp.linspace(0, 2 * jnp.pi * (1 - 1/num_samples), num_samples)
+    xs = jnp.cos(alphas)
+    ys = jnp.sin(alphas)
+    samples = jnp.stack([xs, ys], axis=1)
+    return samples
 
-    def metric_names(self):
-        return ['mean']
+  def metric_names(self):
+    return ['mean']
 
-    def calculate_metrics_batch(self, batch):
-        return vmap(lambda x: jnp.mean(x, axis=0))(batch)[0, 0]
+  def calculate_metrics_batch(self, batch):
+    return vmap(lambda x: jnp.mean(x, axis=0))(batch)[0, 0]
 
-    def get_data_scaler(self, config):
-        def data_scaler(x):
-            return x / jnp.sqrt(2)
-        return data_scaler
+  def get_data_scaler(self, config):
+    def data_scaler(x):
+      return x / jnp.sqrt(2)
+    return data_scaler
 
-    def get_data_inverse_scaler(self, config):
-        def data_inverse_scaler(x):
-            return x * jnp.sqrt(2)
-        return data_inverse_scaler
+  def get_data_inverse_scaler(self, config):
+    def data_inverse_scaler(x):
+      return x * jnp.sqrt(2)
+    return data_inverse_scaler
 
 
 def main(argv):
-    workdir = FLAGS.workdir
-    config = FLAGS.config
-    jax.default_device = jax.devices()[0]
-    # Tip: use CUDA_VISIBLE_DEVICES to restrict the devices visible to jax
-    # ... they must be all the same model of device for pmap to work
-    num_devices =  int(jax.local_device_count()) if config.training.pmap else 1
-    rng = random.PRNGKey(config.seed)
-
-    # Setup SDE
-    if config.training.sde.lower()=='vpsde':
-        sde = sde_lib.VP(beta_min=config.model.beta_min, beta_max=config.model.beta_max)
-    elif config.training.sde.lower()=='vesde':
-        sde = sde_lib.VE(sigma_min=config.model.sigma_min, sigma_max=config.model.sigma_max)
-    else:
-        raise NotImplementedError(f"SDE {config.training.SDE} unknown.")
-
-    # Build data iterators
-    num_samples = 8
-    dataset = CircleDataset(num_samples=num_samples)
-    scaler = dataset.get_data_scaler(config)
-    inverse_scaler = dataset.get_data_inverse_scaler(config)
-    plot_samples(
-        samples=dataset.train_data, index=(0, 1), fname="samples", lims=((-3, 3), (-3, 3)))
-
-    def log_hat_pt(x, t):
-        """
-        Empirical distribution score.
-
-        Args:
-            x: One location in $\mathbb{R}^2$
-            t: time
-        Returns:
-            The empirical log density, as described in the Jupyter notebook
-            .. math::
-                \hat{p}_{t}(x)
-        """
-        mean_coeff = sde.mean_coeff(t)
-        mean = mean_coeff * scaler(dataset.train_data)
-        std = jnp.sqrt(sde.variance(t))
-        potentials = jnp.sum(-(x - mean)**2 / (2 * std**2), axis=1)
-        return logsumexp(potentials, axis=0, b=1/num_samples)
-
-    # Get a jax grad function, which can be batched with vmap
-    nabla_log_hat_pt = jit(vmap(grad(log_hat_pt), in_axes=(0, 0), out_axes=(0)))
-
-    # Running the reverse SDE with the empirical drift
-    plot_score(score=nabla_log_hat_pt, scaler=scaler, t=0.01, area_min=-3, area_max=3, fname="empirical score")
-    outer_solver, inner_solver = get_solver(config, sde, nabla_log_hat_pt)
-    sampler = get_sampler((5760, config.data.image_size),
-                          outer_solver, inner_solver,
-                          denoise=config.sampling.denoise,
-                          stack_samples=False,
-                          inverse_scaler=inverse_scaler)
-    rng, sample_rng = random.split(rng, 2)
-    q_samples, num_function_evaluations = sampler(sample_rng)
-    plot_heatmap(samples=q_samples, area_min=-3, area_max=3, fname="heatmap empirical score")
-
-    # What happens when I perturb the score with a constant?
-    perturbed_score = lambda x, t: nabla_log_hat_pt(x, t) + 1
-    outer_solver, inner_solver = get_solver(config, sde, perturbed_score)
-    sampler = get_sampler((5760, config.data.image_size),
-                          outer_solver, inner_solver,
-                          denoise=config.sampling.denoise,
-                          inverse_scaler=inverse_scaler)
+  workdir = FLAGS.workdir
+  config = FLAGS.config
+  jax.default_device = jax.devices()[0]
+  # Tip: use CUDA_VISIBLE_DEVICES to restrict the devices visible to jax
+  # ... they must be all the same model of device for pmap to work
+  num_devices =  int(jax.local_device_count()) if config.training.pmap else 1
+  rng = random.PRNGKey(config.seed)
+
+  # Setup SDE
+  if config.training.sde.lower()=='vpsde':
+    sde = sde_lib.VP(beta_min=config.model.beta_min, beta_max=config.model.beta_max)
+  elif config.training.sde.lower()=='vesde':
+    sde = sde_lib.VE(sigma_min=config.model.sigma_min, sigma_max=config.model.sigma_max)
+  else:
+    raise NotImplementedError(f"SDE {config.training.SDE} unknown.")
+
+  # Build data iterators
+  num_samples = 8
+  dataset = CircleDataset(num_samples=num_samples)
+  scaler = dataset.get_data_scaler(config)
+  inverse_scaler = dataset.get_data_inverse_scaler(config)
+  plot_samples(
+    samples=dataset.train_data, index=(0, 1), fname="samples", lims=((-3, 3), (-3, 3)))
+
+  def log_hat_pt(x, t):
+    """Empirical distribution score.
+
+    Args:
+      x: One location in $\mathbb{R}^2$
+      t: time
+    Returns:
+      The empirical log density, as described in the Jupyter notebook
+      .. math::
+        \hat{p}_{t}(x)
+    """
+    mean_coeff = sde.mean_coeff(t)
+    mean = mean_coeff * scaler(dataset.train_data)
+    std = jnp.sqrt(sde.variance(t))
+    potentials = jnp.sum(-(x - mean)**2 / (2 * std**2), axis=1)
+    return logsumexp(potentials, axis=0, b=1/num_samples)
+
+  # Get a jax grad function, which can be batched with vmap
+  nabla_log_hat_pt = jit(vmap(grad(log_hat_pt), in_axes=(0, 0), out_axes=(0)))
+
+  # Running the reverse SDE with the empirical drift
+  plot_score(score=nabla_log_hat_pt, scaler=scaler, t=0.01, area_min=-3, area_max=3, fname="empirical score")
+  outer_solver, inner_solver = get_solver(config, sde, nabla_log_hat_pt)
+  sampler = get_sampler((5760, config.data.image_size),
+                        outer_solver, inner_solver,
+                        denoise=config.sampling.denoise,
+                        stack_samples=False,
+                        inverse_scaler=inverse_scaler)
+  rng, sample_rng = random.split(rng, 2)
+  q_samples, _ = sampler(sample_rng)
+  plot_heatmap(samples=q_samples, area_bounds=[-3., 3.], fname="heatmap empirical score")
+
+  # What happens when I perturb the score with a constant?
+  perturbed_score = lambda x, t: nabla_log_hat_pt(x, t) + 1
+  outer_solver, inner_solver = get_solver(config, sde, perturbed_score)
+  sampler = get_sampler((5760, config.data.image_size),
+                        outer_solver, inner_solver,
+                        denoise=config.sampling.denoise,
+                        inverse_scaler=inverse_scaler)
+  rng, sample_rng = random.split(rng, 2)
+  q_samples, _ = sampler(sample_rng)
+  plot_heatmap(samples=q_samples, area_bounds=[-3., 3.], fname="heatmap bounded perturbation")
+
+  if not os.path.exists('/tmp/output0'):
+      time_prev = time.time()
+      params, *_ = train(
+        (config.training.batch_size//jax.local_device_count(), config.data.image_size),
+        config, workdir, dataset)
+      time_delta = time.time() - time_prev
+      print("train time: {}s".format(time_delta))
+
+      # Save params
+      output = serialization.to_bytes(params)
+      f = open('/tmp/output0', 'wb')
+      f.write(output)
+  else:  # Load pre-trained model parameters
+      params = get_model(config).init(
+        rng,
+        jnp.zeros(
+          (config.training.batch_size//jax.local_device_count(), config.data.image_size)
+        ),
+        jnp.ones((config.training.batch_size//jax.local_device_count(),)))
+      f = open('/tmp/output0', 'rb')
+      output = f.read()
+      params = serialization.from_bytes(params, output)
+
+  # Get trained score
+  trained_score = get_score(sde, get_model(config), params, score_scaling=config.training.score_scaling)
+  plot_score(score=trained_score, scaler=scaler, t=0.01, area_min=-3, area_max=3, fname="trained score")
+  outer_solver, inner_solver = get_solver(config, sde, trained_score)
+  sampler = get_sampler(
+    (config.solver.num_outer_steps // num_devices, config.data.image_size),
+    outer_solver,
+    inner_solver,
+    denoise=config.sampling.denoise,
+    inverse_scaler=inverse_scaler)
+
+  if config.training.pmap:
+    sampler = jax.pmap(sampler, axis_name='batch')
+    rng, *sample_rng = random.split(rng, 1 + num_devices)
+    sample_rng = jnp.asarray(sample_rng)
+  else:
     rng, sample_rng = random.split(rng, 2)
-    q_samples, num_function_evaluations = sampler(sample_rng)
-    plot_heatmap(samples=q_samples, area_min=-3, area_max=3, fname="heatmap bounded perturbation")
-
-    if not os.path.exists('/tmp/output0'):
-        time_prev = time.time()
-        params, opt_state, mean_losses = train(
-            (config.training.batch_size//jax.local_device_count(), config.data.image_size),
-            config, workdir, dataset)
-        time_delta = time.time() - time_prev
-        print("train time: {}s".format(time_delta))
-
-        # Save params
-        output = serialization.to_bytes(params)
-        f = open('/tmp/output0', 'wb')
-        f.write(output)
-    else:  # Load pre-trained model parameters
-        params = get_model(config).init(
-            rng,
-            jnp.zeros(
-                (config.training.batch_size//jax.local_device_count(), config.data.image_size)
-            ),
-            jnp.ones((config.training.batch_size//jax.local_device_count(),)))
-        f = open('/tmp/output0', 'rb')
-        output = f.read()
-        params = serialization.from_bytes(params, output)
-
-    # Get trained score
-    trained_score = get_score(sde, get_model(config), params, score_scaling=config.training.score_scaling)
-    plot_score(score=trained_score, scaler=scaler, t=0.01, area_min=-3, area_max=3, fname="trained score")
-    outer_solver, inner_solver = get_solver(config, sde, trained_score)
-    sampler = get_sampler(
-        (config.solver.num_outer_steps // num_devices, config.data.image_size),
-        outer_solver,
-        inner_solver,
-        denoise=config.sampling.denoise,
-        inverse_scaler=inverse_scaler)
-
-    if config.training.pmap:
-        sampler = jax.pmap(sampler, axis_name='batch')
-        rng, *sample_rng = random.split(rng, 1 + num_devices)
-        sample_rng = jnp.asarray(sample_rng)
-    else:
-        rng, sample_rng = random.split(rng, 2)
-
-    q_samples, num_function_evaluations = sampler(sample_rng)
-    q_samples = q_samples.reshape(config.solver.num_outer_steps, config.data.image_size)
-    plot_heatmap(samples=q_samples, area_min=-3, area_max=3, fname="heatmap trained score")
-
-    # Condition on one of the coordinates
-    data = jnp.array([-0.5, 0.0])
-    mask = jnp.array([1, 0])
-    data = jnp.tile(data, (64, 1))
-    data = scaler(data)
-    mask = jnp.tile(mask, (64, 1))
-    inpainter = get_inpainter(
-        outer_solver,
-        stack_samples=False,
-        denoise=config.sampling.denoise,
-        inverse_scaler=inverse_scaler)
-    rng, sample_rng = random.split(rng, 2)
-    q_samples, num_function_evaluations = inpainter(sample_rng, data, mask)
-    plot_heatmap(samples=q_samples, area_min=-3, area_max=3, fname="heatmap inpainted")
+
+  q_samples, _ = sampler(sample_rng)
+  q_samples = q_samples.reshape(config.solver.num_outer_steps, config.data.image_size)
+  plot_heatmap(samples=q_samples, area_bounds=[-3., 3.], fname="heatmap trained score")
+
+  # Condition on one of the coordinates
+  data = jnp.array([-0.5, 0.0])
+  mask = jnp.array([1, 0])
+  data = jnp.tile(data, (64, 1))
+  data = scaler(data)
+  mask = jnp.tile(mask, (64, 1))
+  inpainter = get_inpainter(
+    outer_solver,
+    stack_samples=False,
+    inverse_scaler=inverse_scaler)
+  rng, sample_rng = random.split(rng, 2)
+  q_samples, _ = inpainter(sample_rng, data, mask)
+  plot_heatmap(samples=q_samples, area_bounds=[-3., 3.], fname="heatmap inpainted")
 
 
 if __name__ == "__main__":
-    app.run(main)
+  app.run(main)
diff --git a/examples/example1.py b/examples/example1.py
index d46fb9a..c4c6399 100644
--- a/examples/example1.py
+++ b/examples/example1.py
@@ -1,254 +1,292 @@
-"""Diffusion models introduction.
-
-An example using 1 dimensional image data.
-"""
-import jax
+"""Diffusion models introduction. An example using 1 dimensional image data."""
 from jax import vmap, jit, grad, value_and_grad
 import jax.random as random
 import jax.numpy as jnp
 from jax.scipy.special import logsumexp
 from flax import serialization
-import optax
 from functools import partial
 import matplotlib.pyplot as plt
-from diffusionjax.plot import (
-    plot_score, plot_heatmap, plot_animation)
-from diffusionjax.utils import (
-    get_score, retrain_nn, optimizer, update_step, get_loss,
-    get_inpainter, get_projection_sampler, get_sampler)
+from diffusionjax.plot import plot_score, plot_heatmap, plot_animation
+from diffusionjax.utils import get_score, get_loss, get_sampler
+from diffusionjax.inverse_problems import get_inpainter, get_projection_sampler
 from diffusionjax.solvers import EulerMaruyama
 from diffusionjax.models import MLP
 from diffusionjax.sde import VE
 import numpy as np
 import os
 
-# Dependencies: This example requires mlkernels package,
-# https://github.com/wesselb/mlkernels#installation
+# Dependencies:
+# This example requires mlkernels package, https://github.com/wesselb/mlkernels#installation
 import lab as B
 from mlkernels import Matern52
+# This example requires optax, https://optax.readthedocs.io/en/latest/
+import optax
+
 
 x_max = 5.0
 epsilon = 1e-4
 
 
+#Initialize the optimizer
+optimizer = optax.adam(1e-3)
+
+
+@partial(jit, static_argnums=[4])
+def update_step(params, rng, batch, opt_state, loss):
+  """
+  Takes the gradient of the loss function and updates the model weights (params) using it.
+  Args:
+      params: the current weights of the model
+      rng: random number generator from jax
+      batch: a batch of samples from the training data, representing samples from \mu_text{data}, shape (J, N)
+      opt_state: the internal state of the optimizer
+      loss: A loss function that can be used for score matching training.
+  Returns:
+      The value of the loss function (for metrics), the new params and the new optimizer states function (for metrics),
+      the new params and the new optimizer state.
+  """
+  val, grads = value_and_grad(loss)(params, rng, batch)
+  updates, opt_state = optimizer.update(grads, opt_state)
+  params = optax.apply_updates(params, updates)
+  return val, params, opt_state
+
+
+def retrain_nn(
+    update_step, num_epochs, step_rng, samples, params,
+    opt_state, loss, batch_size=5):
+  train_size = samples.shape[0]
+  batch_size = min(train_size, batch_size)
+  steps_per_epoch = train_size // batch_size
+  mean_losses = jnp.zeros((num_epochs, 1))
+  for i in range(num_epochs):
+    rng, step_rng = random.split(step_rng)
+    perms = random.permutation(step_rng, train_size)
+    perms = perms[:steps_per_epoch * batch_size]  # skip incomplete batch
+    perms = perms.reshape((steps_per_epoch, batch_size))
+    losses = jnp.zeros((jnp.shape(perms)[0], 1))
+    for j, perm in enumerate(perms):
+      batch = samples[perm, :]
+      rng, step_rng = random.split(rng)
+      loss_eval, params, opt_state = update_step(params, step_rng, batch, opt_state, loss)
+      losses = losses.at[j].set(loss_eval)
+    mean_loss = jnp.mean(losses, axis=0)
+    mean_losses = mean_losses.at[i].set(mean_loss)
+    if i % 10 == 0:
+      print("Epoch {:d}, Loss {:.2f} ".format(i, mean_loss[0]))
+  return params, opt_state, mean_losses
+
+
 def sample_image_rgb(rng, num_samples, image_size, kernel, num_channels=1):
-    """Samples from a GMRF
-    """
-    x = np.linspace(-x_max, x_max, image_size)
-    x = x.reshape(image_size, 1)
-    C = B.dense(kernel(x)) + epsilon * B.eye(image_size)
-    u = random.multivariate_normal(rng, mean=jnp.zeros(x.shape[0]), cov=C, shape=(num_samples, num_channels))
-    u = u.transpose((0, 2, 1))
-    return u, C
+  """Samples from a GMRF"""
+  x = np.linspace(-x_max, x_max, image_size)
+  x = x.reshape(image_size, 1)
+  C = B.dense(kernel(x)) + epsilon * B.eye(image_size)
+  u = random.multivariate_normal(rng, mean=jnp.zeros(x.shape[0]), cov=C, shape=(num_samples, num_channels))
+  u = u.transpose((0, 2, 1))
+  return u, C
 
 
 def plot_score_ax_sample(ax, sample, score, t, area_min=-1, area_max=1, fname="plot_score"):
-    @partial(jit, static_argnums=[0,])
-    def helper(score, sample, t, area_min, area_max):
-        x = jnp.linspace(area_min, area_max, 16)
-        x, y = jnp.meshgrid(x, x)
-        grid = jnp.stack([x.flatten(), y.flatten()], axis=1)
-        sample = jnp.tile(sample, (len(x.flatten()), 1, 1, 1))
-        sample.at[:, [0, 1], 0, 0].set(grid)
-        t = jnp.ones((grid.shape[0],)) * t
-        scores = score(sample, t)
-        return grid, scores
-    grid, scores = helper(score, sample, t, area_min, area_max)
-    ax.quiver(grid[:, 0], grid[:, 1], scores[:, 0, 0, 0], scores[:, 1, 0, 0])
+  @partial(jit, static_argnums=[0,])
+  def helper(score, sample, t, area_min, area_max):
+    x = jnp.linspace(area_min, area_max, 16)
+    x, y = jnp.meshgrid(x, x)
+    grid = jnp.stack([x.flatten(), y.flatten()], axis=1)
+    sample = jnp.tile(sample, (len(x.flatten()), 1, 1, 1))
+    sample.at[:, [0, 1], 0, 0].set(grid)
+    t = jnp.ones((grid.shape[0],)) * t
+    scores = score(sample, t)
+    return grid, scores
+  grid, scores = helper(score, sample, t, area_min, area_max)
+  ax.quiver(grid[:, 0], grid[:, 1], scores[:, 0, 0, 0], scores[:, 1, 0, 0])
 
 
 def plot_samples_1D(samples, image_size, fname="samples 1D.png"):
-    x = np.linspace(-x_max, x_max, image_size)
-    plt.plot(x, samples[:, :, 0].T)
-    plt.savefig(fname)
-    plt.close()
+  x = np.linspace(-x_max, x_max, image_size)
+  plt.plot(x, samples[:, :, 0].T)
+  plt.savefig(fname)
+  plt.close()
 
 
 def main():
-    num_epochs = 128
-    rng = random.PRNGKey(2023)
-    rng, step_rng = random.split(rng, 2)
-    num_samples = 18000
-    num_channels = 1
-    image_size = 64  # image size
-
-    samples, C = sample_image_rgb(rng, num_samples=num_samples, image_size=image_size, kernel=Matern52(), num_channels=num_channels)  # (num_samples, image_size, num_channels)
-
-    # Reshape image data
-    samples = samples.reshape(-1, image_size, num_channels)
-    plot_samples_1D(samples[:64], image_size, "samples")
-
-    # Get sde model
-    sde = VE(sigma_min=0.01, sigma_max=3.0)
-
-    def log_hat_pt_tmp(x, t):
-        """
-        Empirical distribution score.
-
-        Args:
-            x: One location in $\mathbb{R}^{image_size}$
-            t: time
-        Returns:
-            The empirical log density, as described in the Jupyter notebook
-            .. math::
-                \hat{p}_{t}(x)
-        """
-        mean, std = sde.marginal_prob(samples[:, [0, 1], 0], t)
-        potentials = jnp.sum(-(x - mean)**2 / (2 * std**2), axis=1)
-        return logsumexp(potentials, axis=0, b=1/num_samples)
-
-    def log_hat_pt(x, t):
-        """
-        Empirical distribution score.
-
-        Args:
-            x: One location in $\mathbb{R}^{image_size}$
-            t: time
-        Returns:
-            The empirical log density, as described in the Jupyter notebook
-            .. math::
-                \hat{p}_{t}(x)
-        """
-        mean, std = sde.marginal_prob(samples, t)
-        losses = -(x - mean)**2 / (2 * std**2)
-        # Needs to be reshaped, since x is an image
-        potentials = jnp.sum(losses.reshape((losses.shape[0], -1)), axis=-1)
-        return logsumexp(potentials, axis=0, b=1/num_samples)
-
-    def nabla_log_pt(x, t):
-        """
-        Args:
-            x: One location in $\mathbb{R}^{image_size}$
-            t: time
-        Returns:
-            The true log density.
-            .. math::
-                p_{t}(x)
-        """
-        x_shape = x.shape
-        v_t = sde.variance(t)
-        m_t = sde.mean_coeff(t)
-        x = x.flatten()
-        score = - jnp.linalg.solve(m_t**2 * C + v_t * jnp.eye(x_shape[0]), x)
-        return score.reshape(x_shape)
-
-    if 0:  # This may take a while
-        # Get a jax grad function, which can be batched with vmap
-        nabla_log_hat_pt_tmp = jit(vmap(grad(log_hat_pt_tmp), in_axes=(0, 0), out_axes=(0)))
-        nabla_log_hat_pt = jit(vmap(grad(log_hat_pt), in_axes=(0, 0), out_axes=(0)))
-
-        # Running the reverse SDE with the empirical score
-        plot_score(score=nabla_log_hat_pt_tmp, scaler=lambda x:x, t=0.01, area_min=-3, area_max=3, fname="empirical score")
-        sampler = get_sampler((64, image_size, num_channels), EulerMaruyama(sde.reverse(nabla_log_hat_pt)))
-        q_samples, num_function_evaluations = sampler(rng)
-        plot_samples_1D(q_samples, image_size=image_size, fname="samples empirical score")
-        plot_heatmap(samples=q_samples[:, [0, 1], 0], area_min=-3, area_max=3, fname="heatmap empirical score")
-
-        # What happens when I perturb the score with a constant?
-        perturbed_score = lambda x, t: nabla_log_hat_pt(x, t) + 10.0 * jnp.ones(jnp.shape(x))
-        rng, step_rng = random.split(rng)
-        sampler = get_sampler((64, image_size, num_channels), EulerMaruyama(sde.reverse(perturbed_score)))
-        q_samples, num_function_evaluations = sampler(rng)
-        plot_samples_1D(q_samples, image_size=image_size, fname="samples bounded perturbation")
-        plot_heatmap(samples=q_samples[:, [0, 1], 0], area_min=-3, area_max=3, fname="heatmap bounded perturbation")
-
-        nabla_log_pt = jit(vmap(nabla_log_pt, in_axes=(0, 0), out_axes=(0)))
-
-        # Running the reverse SDE with the true score
-        sampler = get_sampler((64, image_size, num_channels), EulerMaruyama(sde.reverse(nabla_log_pt)))
-        q_samples, num_function_evaluations = sampler(rng)
-        plot_samples_1D(q_samples, image_size=image_size, fname="samples true score")
-        plot_heatmap(samples=q_samples[:, [0, 1], 0], area_min=-3, area_max=3, fname="heatmap true score")
-
-        # What happens when I perturb the score with a constant?
-        perturbed_score = lambda x, t: nabla_log_pt(x, t) + 10.0 * jnp.ones(jnp.shape(x))
-        sampler = get_sampler((64, image_size, num_channels), EulerMaruyama(sde.reverse(perturbed_score)))
-        rng, sample_rng = random.split(rng, 2)
-        q_samples, num_function_evaluations = sampler(sample_rng)
-        plot_samples_1D(q_samples, image_size=image_size, fname="samples true bounded perturbation")
-        plot_heatmap(samples=q_samples[:, [0, 1], 0], area_min=-3, area_max=3, fname="heatmap true bounded perturbation")
-
-    # Neural network training via score matching
-    batch_size = 64
-    score_model = MLP()
-
-    # Initialize parameters
-    params = score_model.init(step_rng, jnp.zeros((batch_size, image_size, num_channels)), jnp.ones((batch_size,)))
-
-    # Initialize optimizer
-    opt_state = optimizer.init(params)
-
-    if not os.path.exists('/tmp/output1'):
-        solver = EulerMaruyama(sde)
-        # Get loss function
-        loss = get_loss(
-            sde, solver, score_model, score_scaling=True, likelihood_weighting=False,
-            reduce_mean=True, pointwise_t=False)
-
-        # Train with score matching
-        params, opt_state, mean_losses = retrain_nn(
-            update_step=update_step,
-            num_epochs=num_epochs,
-            step_rng=step_rng,
-            samples=samples,
-            params=params,
-            opt_state=opt_state,
-            loss=loss,
-            batch_size=batch_size)
-
-        # Save params
-        output = serialization.to_bytes(params)
-        f = open('/tmp/output1', 'wb')
-        f.write(output)
-    else:  # Load pre-trained model parameters
-        f = open('/tmp/output1', 'rb')
-        output = f.read()
-        params = serialization.from_bytes(params, output)
-
-    # Get trained score
-    trained_score = get_score(sde, score_model, params, score_scaling=True)
-    solver = EulerMaruyama(sde.reverse(trained_score))
-    sampler = get_sampler((512, image_size, num_channels), solver, denoise=True)
+  num_epochs = 128
+  rng = random.PRNGKey(2023)
+  rng, step_rng = random.split(rng, 2)
+  num_samples = 18000
+  num_channels = 1
+  image_size = 64  # image size
+
+  samples, C = sample_image_rgb(rng, num_samples=num_samples, image_size=image_size, kernel=Matern52(), num_channels=num_channels)  # (num_samples, image_size, num_channels)
+
+  # Reshape image data
+  samples = samples.reshape(-1, image_size, num_channels)
+  plot_samples_1D(samples[:64], image_size, "samples")
+
+  # Get sde model
+  sde = VE(sigma_min=0.01, sigma_max=3.0)
+
+  def log_hat_pt_tmp(x, t):
+    """Empirical distribution score.
+
+    Args:
+      x: One location in $\mathbb{R}^{image_size}$
+      t: time
+    Returns:
+      The empirical log density, as described in the Jupyter notebook
+      .. math::
+        \log \hat{p}_{t}(x)
+    """
+    mean, std = sde.marginal_prob(samples[:, [0, 1], 0], t)
+    potentials = jnp.sum(-(x - mean)**2 / (2 * std**2), axis=1)
+    return logsumexp(potentials, axis=0, b=1/num_samples)
 
-    rng, sample_rng = random.split(rng, 2)
-    q_samples, num_function_evaluations = sampler(sample_rng)
-
-    # C_emp = jnp.corrcoef(q_samples[:, :, 0].T)
-    # delta = jnp.linalg.norm(C - C_emp) / image_size
-
-    plot_samples_1D(q_samples[:64], image_size=image_size, fname="samples trained score")
-    plot_heatmap(samples=q_samples[:, [0, 1], 0], area_min=-3, area_max=3, fname="heatmap trained score")
-
-    if 0:
-        frames = 100
-        fig, ax = plt.subplots()
-        def animate(i, ax):
-            ax.clear()
-            plot_score_ax_sample(
-                ax, q_samples[0], trained_score, t=1 - (i / frames), area_min=-5, area_max=5, fname="trained score")
-        # Plot animation of the trained score over time
-        plot_animation(fig, ax, animate, frames, "trained_score")
-
-    # Condition on one of the coordinates
-    data = jnp.zeros((image_size, num_channels))
-    data = data.at[[0, -1], 0].set([-1.0, 1.0])
-    mask = jnp.zeros((image_size, num_channels), dtype=int)
-    mask = mask.at[[0, -1], 0].set([1, 1])
-    data = jnp.tile(data, (5, 1, 1))
-    mask = jnp.tile(mask, (5, 1, 1))
-
-    # Get inpainter
-    inpainter = get_inpainter(solver, stack_samples=False)
-    rng, sample_rng = random.split(rng, 2)
-    q_samples, num_function_evaluations = inpainter(sample_rng, data, mask)
-    plot_samples_1D(q_samples, image_size=image_size, fname="samples inpainted")
+  def log_hat_pt(x, t):
+    """Empirical distribution score.
 
-    # Get projection sampler
-    projection_sampler = get_projection_sampler(solver, stack_samples=False)
+    Returns:
+      The empirical log density, as described in the Jupyter notebook
+      .. math::
+        \log \hat{p}_{t}(x)
+    """
+    mean, std = sde.marginal_prob(samples, t)
+    losses = -(x - mean)**2 / (2 * std**2)
+    # Needs to be reshaped, since x is an image
+    potentials = jnp.sum(losses.reshape((losses.shape[0], -1)), axis=-1)
+    return logsumexp(potentials, axis=0, b=1/num_samples)
+
+  def nabla_log_pt(x, t):
+    """Score.
+
+    Returns:
+      The true log density.
+      .. math::
+        \nabla_{x} \log p_{t}(x)
+    """
+    x_shape = x.shape
+    v_t = sde.variance(t)
+    m_t = sde.mean_coeff(t)
+    x = x.flatten()
+    score = - jnp.linalg.solve(m_t**2 * C + v_t * jnp.eye(x_shape[0]), x)
+    return score.reshape(x_shape)
+
+  if 0:  # This may take a while
+    # Get a jax grad function, which can be batched with vmap
+    nabla_log_hat_pt_tmp = jit(vmap(grad(log_hat_pt_tmp), in_axes=(0, 0), out_axes=(0)))
+    nabla_log_hat_pt = jit(vmap(grad(log_hat_pt), in_axes=(0, 0), out_axes=(0)))
+
+    # Running the reverse SDE with the empirical score
+    plot_score(score=nabla_log_hat_pt_tmp, scaler=lambda x:x, t=0.01, area_min=-3, area_max=3, fname="empirical score")
+    sampler = get_sampler((64, image_size, num_channels), EulerMaruyama(sde.reverse(nabla_log_hat_pt)))
+    q_samples, _ = sampler(rng)
+    plot_samples_1D(q_samples, image_size=image_size, fname="samples empirical score")
+    plot_heatmap(samples=q_samples[:, [0, 1], 0], area_bounds=[-3., 3.], fname="heatmap empirical score")
+
+    # What happens when I perturb the score with a constant?
+    perturbed_score = lambda x, t: nabla_log_hat_pt(x, t) + 10.0 * jnp.ones(jnp.shape(x))
+    rng, step_rng = random.split(rng)
+    sampler = get_sampler((64, image_size, num_channels), EulerMaruyama(sde.reverse(perturbed_score)))
+    q_samples, _ = sampler(rng)
+    plot_samples_1D(q_samples, image_size=image_size, fname="samples bounded perturbation")
+    plot_heatmap(samples=q_samples[:, [0, 1], 0], area_bounds=[-3., 3.], fname="heatmap bounded perturbation")
+
+    nabla_log_pt = jit(vmap(nabla_log_pt, in_axes=(0, 0), out_axes=(0)))
+
+    # Running the reverse SDE with the true score
+    sampler = get_sampler((64, image_size, num_channels), EulerMaruyama(sde.reverse(nabla_log_pt)))
+    q_samples, num_function_evaluations = sampler(rng)
+    plot_samples_1D(q_samples, image_size=image_size, fname="samples true score")
+    plot_heatmap(samples=q_samples[:, [0, 1], 0], area_bounds=[-3., 3.], fname="heatmap true score")
+
+    # What happens when I perturb the score with a constant?
+    perturbed_score = lambda x, t: nabla_log_pt(x, t) + 10.0 * jnp.ones(jnp.shape(x))
+    sampler = get_sampler((64, image_size, num_channels), EulerMaruyama(sde.reverse(perturbed_score)))
     rng, sample_rng = random.split(rng, 2)
-    q_samples, num_function_evaluations = projection_sampler(sample_rng, data, mask, 1e-2)
-    plot_samples_1D(q_samples, image_size=image_size, fname="samples projected")
+    q_samples, _ = sampler(sample_rng)
+    plot_samples_1D(q_samples, image_size=image_size, fname="samples true bounded perturbation")
+    plot_heatmap(samples=q_samples[:, [0, 1], 0], area_bounds=[-3., 3.], fname="heatmap true bounded perturbation")
+
+  # Neural network training via score matching
+  batch_size = 64
+  score_model = MLP()
+
+  # Initialize parameters
+  params = score_model.init(step_rng, jnp.zeros((batch_size, image_size, num_channels)), jnp.ones((batch_size,)))
+
+  # Initialize optimizer
+  opt_state = optimizer.init(params)
+
+  if not os.path.exists('/tmp/output1'):
+    solver = EulerMaruyama(sde)
+
+    # Get loss function
+    loss = get_loss(
+      sde, solver, score_model, score_scaling=True, likelihood_weighting=False,
+      reduce_mean=True, pointwise_t=False)
+
+    # Train with score matching
+    params, opt_state, _ = retrain_nn(
+      update_step=update_step,
+      num_epochs=num_epochs,
+      step_rng=step_rng,
+      samples=samples,
+      params=params,
+      opt_state=opt_state,
+      loss=loss,
+      batch_size=batch_size)
+
+    # Save params
+    output = serialization.to_bytes(params)
+    f = open('/tmp/output1', 'wb')
+    f.write(output)
+  else:  # Load pre-trained model parameters
+    f = open('/tmp/output1', 'rb')
+    output = f.read()
+    params = serialization.from_bytes(params, output)
+
+  # Get trained score
+  trained_score = get_score(sde, score_model, params, score_scaling=True)
+  solver = EulerMaruyama(sde.reverse(trained_score))
+  sampler = get_sampler((512, image_size, num_channels), solver, denoise=True)
+
+  rng, sample_rng = random.split(rng, 2)
+  q_samples, num_function_evaluations = sampler(sample_rng)
+
+  # C_emp = jnp.corrcoef(q_samples[:, :, 0].T)
+  # delta = jnp.linalg.norm(C - C_emp) / image_size
+
+  plot_samples_1D(q_samples[:64], image_size=image_size, fname="samples trained score")
+  plot_heatmap(samples=q_samples[:, [0, 1], 0], area_bounds=[-3., 3.], fname="heatmap trained score")
+
+  if 0:
+    frames = 100
+    fig, ax = plt.subplots()
+    def animate(i, ax):
+      ax.clear()
+      plot_score_ax_sample(
+        ax, q_samples[0], trained_score, t=1 - (i / frames), area_min=-5, area_max=5, fname="trained score")
+    # Plot animation of the trained score over time
+    plot_animation(fig, ax, animate, frames, "trained_score")
+
+  # Condition on one of the coordinates
+  data = jnp.zeros((image_size, num_channels))
+  data = data.at[[0, -1], 0].set([-1.0, 1.0])
+  mask = jnp.zeros((image_size, num_channels), dtype=int)
+  mask = mask.at[[0, -1], 0].set([1, 1])
+  data = jnp.tile(data, (5, 1, 1))
+  mask = jnp.tile(mask, (5, 1, 1))
+
+  # Get inpainter
+  inpainter = get_inpainter(solver, stack_samples=False)
+  rng, sample_rng = random.split(rng, 2)
+  q_samples, _ = inpainter(sample_rng, data, mask)
+  plot_samples_1D(q_samples, image_size=image_size, fname="samples inpainted")
+
+  # Get projection sampler
+  projection_sampler = get_projection_sampler(solver, stack_samples=False)
+  rng, sample_rng = random.split(rng, 2)
+  q_samples, _ = projection_sampler(sample_rng, data, mask, 1e-2)
+  plot_samples_1D(q_samples, image_size=image_size, fname="samples projected")
 
 
 if __name__ == "__main__":
-    main()
+  main()
diff --git a/examples/example2.py b/examples/example2.py
index 86609b0..f0b70a6 100644
--- a/examples/example2.py
+++ b/examples/example2.py
@@ -1,203 +1,232 @@
-"""Diffusion models introduction.
-
-An example using 2 dimensional image data.
-"""
+"""Diffusion models introduction. An example using 2 dimensional image data."""
 import jax
-from jax import vmap, jit, grad
-from functools import partial
+from jax import vmap, jit, grad, value_and_grad
 import jax.random as random
 import jax.numpy as jnp
 from jax.scipy.special import logsumexp
-import optax
-from functools import partial
 from flax import serialization
+from functools import partial
 import matplotlib.pyplot as plt
 from diffusionjax.plot import plot_samples, plot_heatmap
-from diffusionjax.utils import (
-    get_score, retrain_nn, optimizer, update_step, get_loss,
-    get_sampler)
+from diffusionjax.utils import get_score, get_loss, get_sampler
 from diffusionjax.solvers import EulerMaruyama, Annealed
 from diffusionjax.models import CNN
 from diffusionjax.sde import VP, UDLangevin
 import numpy as np
 import os
 
-# Dependencies: This example requires mlkernels package,
-# https://github.com/wesselb/mlkernels#installation
+# Dependencies:
+# This example requires mlkernels package, https://github.com/wesselb/mlkernels#installation
 from mlkernels import Matern52
 import lab as B
+# This example requires optax, https://optax.readthedocs.io/en/latest/
+import optax
+
 
 x_max = 5.0
 epsilon = 1e-4
 
 
+#Initialize the optimizer
+optimizer = optax.adam(1e-3)
+
+
+@partial(jit, static_argnums=[4])
+def update_step(params, rng, batch, opt_state, loss):
+  """
+  Takes the gradient of the loss function and updates the model weights (params) using it.
+  Args:
+      params: the current weights of the model
+      rng: random number generator from jax
+      batch: a batch of samples from the training data, representing samples from \mu_text{data}, shape (J, N)
+      opt_state: the internal state of the optimizer
+      loss: A loss function that can be used for score matching training.
+  Returns:
+      The value of the loss function (for metrics), the new params and the new optimizer states function (for metrics),
+      the new params and the new optimizer state.
+  """
+  val, grads = value_and_grad(loss)(params, rng, batch)
+  updates, opt_state = optimizer.update(grads, opt_state)
+  params = optax.apply_updates(params, updates)
+  return val, params, opt_state
+
+
+def retrain_nn(
+    update_step, num_epochs, step_rng, samples, params,
+    opt_state, loss, batch_size=5):
+  train_size = samples.shape[0]
+  batch_size = min(train_size, batch_size)
+  steps_per_epoch = train_size // batch_size
+  mean_losses = jnp.zeros((num_epochs, 1))
+  for i in range(num_epochs):
+    rng, step_rng = random.split(step_rng)
+    perms = random.permutation(step_rng, train_size)
+    perms = perms[:steps_per_epoch * batch_size]  # skip incomplete batch
+    perms = perms.reshape((steps_per_epoch, batch_size))
+    losses = jnp.zeros((jnp.shape(perms)[0], 1))
+    for j, perm in enumerate(perms):
+      batch = samples[perm, :]
+      rng, step_rng = random.split(rng)
+      loss_eval, params, opt_state = update_step(params, step_rng, batch, opt_state, loss)
+      losses = losses.at[j].set(loss_eval)
+    mean_loss = jnp.mean(losses, axis=0)
+    mean_losses = mean_losses.at[i].set(mean_loss)
+    if i % 10 == 0:
+      print("Epoch {:d}, Loss {:.2f} ".format(i, mean_loss[0]))
+  return params, opt_state, mean_losses
+
+
 def image_grid(x, image_size, num_channels):
-    img = x.reshape(-1, image_size, image_size, num_channels)
-    w = int(np.sqrt(img.shape[0]))
-    return img.reshape((w, w, image_size, image_size, num_channels)).transpose((0, 2, 1, 3, 4)).reshape((w * image_size, w * image_size, num_channels))
+  img = x.reshape(-1, image_size, image_size, num_channels)
+  w = int(np.sqrt(img.shape[0]))
+  return img.reshape((w, w, image_size, image_size, num_channels)).transpose((0, 2, 1, 3, 4)).reshape((w * image_size, w * image_size, num_channels))
 
 
 def plot_samples(x, image_size=32, num_channels=3, fname="samples"):
-    img = image_grid(x, image_size, num_channels)
-    plt.figure(figsize=(8,8))
-    plt.axis('off')
-    plt.imshow(img)
-    plt.savefig(fname)
-    plt.close()
+  img = image_grid(x, image_size, num_channels)
+  plt.figure(figsize=(8,8))
+  plt.axis('off')
+  plt.imshow(img)
+  plt.savefig(fname)
+  plt.close()
 
 
 def sample_image_rgb(rng, num_samples, image_size, kernel, num_channels):
-    """Samples from a GMRF."""
-    x = np.linspace(-x_max, x_max, image_size)
-    y = np.linspace(-x_max, x_max, image_size)
-    xx, yy = np.meshgrid(x, y)
-    xx = xx.reshape(image_size**2, 1)
-    yy = yy.reshape(image_size**2, 1)
-    z = np.hstack((xx, yy))
-    C = B.dense(kernel(z))  + epsilon * B.eye(image_size**2)
-    u = random.multivariate_normal(rng, mean=jnp.zeros(xx.shape[0]), cov=C, shape=(num_samples, num_channels))
-    u = u.transpose((0, 2, 1))
-    return u, C
+  """Samples from a GMRF."""
+  x = np.linspace(-x_max, x_max, image_size)
+  y = np.linspace(-x_max, x_max, image_size)
+  xx, yy = np.meshgrid(x, y)
+  xx = xx.reshape(image_size**2, 1)
+  yy = yy.reshape(image_size**2, 1)
+  z = np.hstack((xx, yy))
+  C = B.dense(kernel(z))  + epsilon * B.eye(image_size**2)
+  u = random.multivariate_normal(rng, mean=jnp.zeros(xx.shape[0]), cov=C, shape=(num_samples, num_channels))
+  u = u.transpose((0, 2, 1))
+  return u, C
 
 
 def plot_samples_1D(samples, image_size, fname="samples 1D.png"):
-    x = np.linspace(-x_max, x_max, image_size)
-    plt.plot(x, samples[:, :, 0, 0].T)
-    plt.savefig(fname)
-    plt.close()
+  x = np.linspace(-x_max, x_max, image_size)
+  plt.plot(x, samples[:, :, 0, 0].T)
+  plt.savefig(fname)
+  plt.close()
 
 
 def main():
-    num_epochs = 128
-    rng = random.PRNGKey(2023)
-    rng, step_rng = random.split(rng, 2)
-    num_samples = 144
-    num_channels = 1
-    image_size = 32  # image size
-    num_steps = 1000
-
-    samples, C = sample_image_rgb(rng, num_samples=num_samples, image_size=image_size, kernel=Matern52(), num_channels=num_channels)  # (num_samples, image_size**2, num_channels)
-    plot_samples(samples[:64], image_size=image_size, num_channels=num_channels)
-    # Reshape image data
-    samples = samples.reshape(-1, image_size, image_size, num_channels)
-    plot_samples_1D(samples[:64], image_size, "samples 1D")
-
-    # Get sde model, variance preserving (VP) a.k.a. time-changed Ohrnstein Uhlenbeck (OU)
-    sde = VP(beta_min=0.1, beta_max=25.0)
-
-    def log_hat_pt(x, t):
-        """
-        Empirical distribution score.
-
-        Args:
-            x: One location in $\mathbb{R}^{image_size}$
-            t: time
-        Returns:
-            The empirical log density, as described in the Jupyter notebook
-            .. math::
-                \hat{p}_{t}(x)
-        """
-        mean, std = sde.marginal_prob(samples, t)
-        losses = -(x - mean)**2 / (2 * std**2)
-        # Needs to be reshaped, since x is an image
-        potentials = jnp.sum(losses.reshape((losses.shape[0], -1)), axis=-1)
-        return logsumexp(potentials, axis=0, b=1/num_samples)
-
-    def nabla_log_pt(x, t):
-        """
-        Args:
-            x: One location in $\mathbb{R}^{image_size}$
-            t: time
-        Returns:
-            The true log density.
-            .. math::
-                p_{t}(x)
-        """
-        x_shape = x.shape
-        v_t = sde.variance(t)
-        m_t = sde.mean_coeff(t)
-        x = x.flatten()
-        score = -jnp.linalg.solve(m_t**2 * C + v_t * jnp.eye(x_shape[0] * x_shape[1]), x)
-        return score.reshape(x_shape)
-
-    if 0:  # this may take a while
-        # Get a jax grad function, which can be batched with vmap
-        nabla_log_hat_pt = jit(vmap(grad(log_hat_pt), in_axes=(0, 0), out_axes=(0)))
-        # nabla_log_pt = jit(vmap(nabla_log_pt, in_axes=(0, 0), out_axes=(0)))
-        # Running the reverse SDE with the empirical score
-        sampler = jax.pmap(get_sampler((64, image_size, image_size, num_channels), EulerMaruyama(sde.reverse(nabla_log_hat_pt), num_steps=num_steps)), axis_name='batch')
-        q_samples, num_function_evaluations = sampler(rng)
-        plot_samples(q_samples, image_size=image_size, num_channels=num_channels, fname="samples empirical score")
-        plot_samples_1D(q_samples, image_size, "samples 1D empirical score")
-        plot_heatmap(samples=q_samples[:, [0, 1], 0, 0], area_min=-3, area_max=3, fname="heatmap empirical score")
-        # What happens when I perturb the score with a constant?
-        perturbed_score = lambda x, t: nabla_log_hat_pt(x, t) + 10.0 * jnp.ones(jnp.shape(x))
-        rng, step_rng = random.split(rng)
-        sampler = jax.pmap(get_sampler((64, image_size, image_size, num_channels), EulerMaruyama(sde.reverse(perturbed_score), num_steps=num_steps)), axis_name='batch')
-        q_samples, num_function_evaluations = sampler(rng)
-        plot_samples(q_samples, image_size=image_size, num_channels=num_channels, fname="samples bounded perturbation")
-        plot_heatmap(samples=q_samples[:, [0, 1], 0, 0], area_min=-3, area_max=3, fname="heatmap bounded perturbation")
-
-    # Neural network training via score matching
-    batch_size = 16
-    score_model = CNN()
-    # Initialize parameters
-    params = score_model.init(step_rng, jnp.zeros((batch_size, image_size, image_size, num_channels)), jnp.ones((batch_size,)))
-    # Initialize optimizer
-    opt_state = optimizer.init(params)
-
-    if not os.path.exists('/tmp/output2'):
-        # Get loss function
-        solver = EulerMaruyama(sde, num_steps=num_steps)
-        loss = get_loss(
-            sde, solver, score_model, score_scaling=True, likelihood_weighting=False,
-            reduce_mean=True, pointwise_t=False)
-        # Train with score matching
-        params, opt_state, mean_losses = retrain_nn(
-            update_step=update_step,
-            num_epochs=num_epochs,
-            step_rng=step_rng,
-            samples=samples,
-            params=params,
-            opt_state=opt_state,
-            loss=loss,
-            batch_size=batch_size)
-
-        # Save params
-        output = serialization.to_bytes(params)
-        f = open('/tmp/output2', 'wb')
-        f.write(output)
-    else:  # Load pre-trained model parameters
-        f = open('/tmp/output2', 'rb')
-        output = f.read()
-        params = serialization.from_bytes(params, output)
-
-    # Get trained score
-    trained_score = get_score(sde, score_model, params, score_scaling=True)
-
-    # Get the outer loop of a numerical solver, also known as "predictor"
-    outer_solver = EulerMaruyama(sde.reverse(trained_score), num_steps=num_steps)
-
-    # Get the inner loop of a numerical solver, also known as "corrector"
-    inner_solver = Annealed(sde.corrector(UDLangevin, trained_score), num_steps=2, snr=0.01)
-
-    # pmap across devices. pmap assumes devices are identical model. If this is not the case,
-    # use the devices argument in pmap
-    num_devices =  jax.local_device_count()
-    sampler = jax.pmap(
-        get_sampler((64//num_devices, image_size, image_size, num_channels), outer_solver, inner_solver, denoise=True),
-        axis_name='batch',
-        # devices = jax.devices()[:],
-    )
-
-    rng, *sample_rng = random.split(rng, num_devices + 1)
-    sample_rng = jnp.asarray(sample_rng)
-    q_samples, num_function_evaluations = sampler(sample_rng)
-    q_samples = q_samples.reshape(64, image_size, image_size, num_channels)
-    plot_samples(q_samples, image_size=image_size, num_channels=num_channels, fname="samples trained score")
-    plot_samples_1D(q_samples, image_size, fname="samples 1D trained score")
-    plot_heatmap(samples=q_samples[:, [0, 1], 0, 0], area_min=-3, area_max=3, fname="heatmap trained score")
+  num_epochs = 128
+  rng = random.PRNGKey(2023)
+  rng, step_rng = random.split(rng, 2)
+  num_samples = 144
+  num_channels = 1
+  image_size = 32  # image size
+  num_steps = 1000
+
+  # Get and handle image data
+  samples, _ = sample_image_rgb(rng, num_samples=num_samples, image_size=image_size, kernel=Matern52(), num_channels=num_channels)  # (num_samples, image_size**2, num_channels)
+  plot_samples(samples[:64], image_size=image_size, num_channels=num_channels)
+  samples = samples.reshape(-1, image_size, image_size, num_channels)
+  plot_samples_1D(samples[:64], image_size, "samples 1D")
+
+  # Get sde model, variance preserving (VP) a.k.a. time-changed Ohrnstein Uhlenbeck (OU)
+  sde = VP(beta_min=0.1, beta_max=25.0)
+
+  def log_hat_pt(x, t):
+    """
+    Empirical distribution score.
+
+    Args:
+      x: One location in $\mathbb{R}^{image_size}$
+      t: time
+    Returns:
+      The empirical log density, as described in the Jupyter notebook
+      .. math::
+        \hat{p}_{t}(x)
+    """
+    mean, std = sde.marginal_prob(samples, t)
+    losses = -(x - mean)**2 / (2 * std**2)
+    # Needs to be reshaped, since x is an image
+    potentials = jnp.sum(losses.reshape((losses.shape[0], -1)), axis=-1)
+    return logsumexp(potentials, axis=0, b=1/num_samples)
+
+  if 0:  # this may take a while
+    # Get a jax grad function, which can be batched with vmap
+    nabla_log_hat_pt = jit(vmap(grad(log_hat_pt), in_axes=(0, 0), out_axes=(0)))
+    # Running the reverse SDE with the empirical score
+    sampler = jax.pmap(get_sampler((64, image_size, image_size, num_channels), EulerMaruyama(sde.reverse(nabla_log_hat_pt), num_steps=num_steps)), axis_name='batch')
+    q_samples, _ = sampler(rng)
+    plot_samples(q_samples, image_size=image_size, num_channels=num_channels, fname="samples empirical score")
+    plot_samples_1D(q_samples, image_size, "samples 1D empirical score")
+    plot_heatmap(samples=q_samples[:, [0, 1], 0, 0], area_min=-3, area_max=3, fname="heatmap empirical score")
+    # What happens when I perturb the score with a constant?
+    perturbed_score = lambda x, t: nabla_log_hat_pt(x, t) + 10.0 * jnp.ones(jnp.shape(x))
+    rng, step_rng = random.split(rng)
+    sampler = jax.pmap(get_sampler((64, image_size, image_size, num_channels), EulerMaruyama(sde.reverse(perturbed_score), num_steps=num_steps)), axis_name='batch')
+    q_samples, num_function_evaluations = sampler(rng)
+    plot_samples(q_samples, image_size=image_size, num_channels=num_channels, fname="samples bounded perturbation")
+    plot_heatmap(samples=q_samples[:, [0, 1], 0, 0], area_min=-3, area_max=3, fname="heatmap bounded perturbation")
+
+  # Neural network training via score matching
+  batch_size = 16
+  score_model = CNN()
+
+  # Initialize parameters
+  params = score_model.init(step_rng, jnp.zeros((batch_size, image_size, image_size, num_channels)), jnp.ones((batch_size,)))
+
+  # Initialize optimizer
+  opt_state = optimizer.init(params)
+
+  if not os.path.exists('/tmp/output2'):
+    # Get loss function
+    solver = EulerMaruyama(sde, num_steps=num_steps)
+    loss = get_loss(
+      sde, solver, score_model, score_scaling=True, likelihood_weighting=False,
+      reduce_mean=True, pointwise_t=False)
+
+    # Train with score matching
+    params, opt_state, _ = retrain_nn(
+      update_step=update_step,
+      num_epochs=num_epochs,
+      step_rng=step_rng,
+      samples=samples,
+      params=params,
+      opt_state=opt_state,
+      loss=loss,
+      batch_size=batch_size)
+
+    # Save params
+    output = serialization.to_bytes(params)
+    f = open('/tmp/output2', 'wb')
+    f.write(output)
+  else:  # Load pre-trained model parameters
+    f = open('/tmp/output2', 'rb')
+    output = f.read()
+    params = serialization.from_bytes(params, output)
+
+  # Get trained score
+  trained_score = get_score(sde, score_model, params, score_scaling=True)
+
+  # Get the outer loop of a numerical solver, also known as "predictor"
+  outer_solver = EulerMaruyama(sde.reverse(trained_score), num_steps=num_steps)
+
+  # Get the inner loop of a numerical solver, also known as "corrector"
+  inner_solver = Annealed(sde.corrector(UDLangevin, trained_score), num_steps=2, snr=0.01)
+
+  # pmap across devices. pmap assumes devices are identical model. If this is not the case,
+  # use the devices argument in pmap
+  num_devices =  jax.local_device_count()
+  sampler = jax.pmap(
+    get_sampler((64//num_devices, image_size, image_size, num_channels), outer_solver, inner_solver, denoise=True),
+    axis_name='batch',
+    # devices = jax.devices()[:],
+  )
+  rng, *sample_rng = random.split(rng, num_devices + 1)
+  sample_rng = jnp.asarray(sample_rng)
+  q_samples, _ = sampler(sample_rng)
+  q_samples = q_samples.reshape(64, image_size, image_size, num_channels)
+  plot_samples(q_samples, image_size=image_size, num_channels=num_channels, fname="samples trained score")
+  plot_samples_1D(q_samples, image_size, fname="samples 1D trained score")
+  plot_heatmap(samples=q_samples[:, [0, 1], 0, 0], area_bounds=[-3., 3.], fname="heatmap trained score")
 
 
 if __name__ == "__main__":
-    main()
+  main()
diff --git a/pyproject.toml b/pyproject.toml
index 9bf1ecf..fce4ade 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,10 +1,10 @@
 [build-system]
 requires = [
-    "setuptools>=50.0",
-    "setuptools_scm[toml]>=6.0",
-    "setuptools_scm_git_archive",
-    "wheel>=0.33",
-    "numpy>=1.16",
+  "setuptools>=50.0",
+  "setuptools_scm[toml]>=6.0",
+  "setuptools_scm_git_archive",
+  "wheel>=0.33",
+  "numpy>=1.16",
 ]
 
 [tool.setuptools_scm]
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index 38c18c9..0000000
--- a/requirements.txt
+++ /dev/null
@@ -1,9 +0,0 @@
-# Requirements for development, testing, and documentation.
--e .
-pytest
-pytest-cov
-coveralls
-jax>=0.4.1
-jaxlib>=0.4.1
-setuptools_scm[toml]
-setuptools_scm_git_archive
diff --git a/setup.py b/setup.py
index 44eefcb..247158b 100644
--- a/setup.py
+++ b/setup.py
@@ -6,7 +6,6 @@
 will occur.
 """
 from setuptools import setup, find_packages
-import subprocess
 import pathlib
 
 
@@ -20,27 +19,51 @@
 LICENSE = (HERE / "LICENSE.rst").read_text()
 
 setup(
-    name="diffusionjax",
-    # python_requires=">=3.8",
-    description="diffusionjax is a simple and accessible diffusion models package in JAX",
-    long_description=README,
-    long_description_content_type="text/markdown",
-    url="https://github.com/bb515/diffusionjax",
-    author="Jakiw Pidstrigach and Benjamin Boys",
-    license="MIT",
-    license_file=LICENSE,
-    packages=find_packages(exclude=['*.test']),
-    include_package_data=True,
-    install_requires=[
-        'numpy',
-        'scipy',
-        'matplotlib',
-        'seaborn',
-        'optax',
-        'flax',
-        'ml_collections',
-        'tqdm',
-        'absl-py',
-        'wandb'
-        ]
-    )
+  name="diffusionjax",
+  # python_requires=">=3.8",
+  description="diffusionjax is a simple and accessible diffusion models package in JAX",
+  long_description=README,
+  long_description_content_type="text/markdown",
+  url="https://github.com/bb515/diffusionjax",
+  author="Benjamin Boys and Jakiw Pidstrigach",
+  license="MIT",
+  license_file=LICENSE,
+  packages=find_packages(exclude=["*.test"]),
+  install_requires=[
+    "numpy",
+    "scipy",
+    "matplotlib",
+    "flax",
+    "ml_collections",
+    "tqdm",
+    "absl-py",
+    "wandb",
+    ],
+  extras_require={
+    'linting': [
+      "flake8",
+      "pylint",
+      "mypy",
+      "typing-extensions",
+      "pre-commit",
+      "ruff",
+      'jaxtyping',
+    ],
+    'testing': [
+      "pytest",
+      "pytest-xdist",
+      "pytest-cov",
+      "coveralls",
+      "jax>=0.4.1",
+      "jaxlib>=0.4.1",
+      "setuptools_scm[toml]",
+      "setuptools_scm_git_archive",
+    ],
+    'examples': [
+      "optax",
+      "orbax-checkpoint",
+      "torch",
+      "mlkernels",
+    ],
+  },
+  include_package_data=True)


Test patch:
diff --git a/diffusionjax/test/test_utils.py b/diffusionjax/test/test_utils.py
index 6335771..9f39d86 100644
--- a/diffusionjax/test/test_utils.py
+++ b/diffusionjax/test/test_utils.py
@@ -5,9 +5,9 @@
 
 def test_batch_mul():
     """Placeholder test for `:meth:batch_mul` to test CI"""
-    a = jnp.ones((2,))
+    a = jnp.ones((2,)) * 2.
     bs = [jnp.zeros((2,)), jnp.ones((2,)), jnp.ones((2,)) * jnp.pi]
-    c_expecteds = [jnp.zeros((2,)), jnp.ones((2,)), jnp.ones((2,)) * jnp.pi]
+    c_expecteds = [jnp.zeros((2,)), 2. * jnp.ones((2,)), 2. * jnp.ones((2,)) * jnp.pi]
     for i, b in enumerate(bs):
         c = batch_mul(a, b)
         assert jnp.allclose(c, c_expecteds[i])