Repository: pyro-ppl/pyro

Issue:
Pyro NUTS inference slow when compared to the same PyMC3 model
For the same models below, the PyMC3 model finishes within a second, whereas the Pyro model has an extremely slow rate of progress. I came across this perf issue while working on a larger model, and isolated it to the example below. I am not sure if this is a perf bug, or whether PyMC3 uses some adaptation tricks, but it is worth looking into and learning from. The cause of the slowdown is the very small step size in Pyro (~10^-9), whereas in PyMC3, the step size throughout is of the order of ~ 10^-3. This is also not a case of bad initialization, the same can be observed for different random seeds.

Also refer to #1470. 

```python
import pymc3 as pm
import pyro

import pyro.distributions as dist
from pyro.infer.mcmc import MCMC, NUTS

import torch


def pm_model(data):
    with pm.Model() as model:
        alpha = pm.HalfCauchy('alpha', beta=2)
        beta = pm.HalfCauchy('beta', beta=2)
        l = pm.Gamma('l', alpha, beta)
        pm.Poisson('obs', l, observed=data)
    return model


def pyro_model(data):
    alpha = pyro.sample('alpha', dist.HalfCauchy(2.))
    beta = pyro.sample('beta', dist.HalfCauchy(2.))
    l = pyro.sample('l', dist.Gamma(alpha, beta))
    pyro.sample('obs', dist.Poisson(l), obs=data)


def get_samples_pymc(data, num_samples=200, warmup_steps=200):
    data = data.numpy()
    with pm_model(data):
        trace = pm.sample(draws=num_samples, tune=warmup_steps, chains=1)


def get_samples_pyro(data, num_samples=200, warmup_steps=200):
    data = data
    nuts_kernel = NUTS(pyro_model,
                       adapt_step_size=True,
                       adapt_mass_matrix=True,
                       jit_compile=True,
                       full_mass=True)
    mcmc_run = MCMC(nuts_kernel,
                    num_samples=num_samples,
                    warmup_steps=warmup_steps,
                    num_chains=1).run(data)


data = torch.tensor(4805497.)
get_samples_pymc(data)  # finishes in a second.
get_samples_pyro(data)  # orders of magnitude slower.
```



Discussion:
Here are some of my thoughts on this:
- I think that this model is not a good setup for the data. The data here is pretty large for a Poisson distribution, so the initial energy will be pretty large. Hence it makes sense to have a very small initial step_size (given that we fix inverse_mass_matrix during the initial window adaptation).
- I guess we can fix this problem by initializing the initial inverse mass matrix to a small value or turning off adapt_step_size, and set step_size = 1e-4 for example.
- PyMC3 does not use window adaptation, so somehow inverse mass matrix will help step_size not move to a very small value.

Anyway, the model setup is not good, so the slowness is acceptable in my opinion.
> I think that this model is not a good setup for the data.

This is from a larger model from a user (originally written in PyMC3), and the above is just a minimal example that is representative of the issue. I do agree that this is not an ideal model setup, but I also think that we can use this simple case to make NUTS more robust (and awesome), both because bigger more complex examples might exhibit similar patterns, and for a better user experience. Consider for example a model that worked well on smaller observations, but became suddenly slow because the user got some new data that was way larger. In such cases, we aren't providing any warning, but effectively crawling to a halt. Both PyMC3 and Stan (which uses the same windowed adaptation scheme) are extremely fast on this, though Stan throws a few warnings about divergent transitions.

> I guess we can fix this problem by initializing the initial inverse mass matrix to a small value or turning off adapt_step_size, and set step_size = 1e-4 for example.

Ideally, we would like to aim for inference to be as automatic as possible, or at least raise warnings so that users know what to do when data/model seems aberrant. Setting the step size upfront is not an option because the actual model is much larger, and it may not work for the remaining model. Likewise, it would be much better for us to start with reasonable mass matrix (somehow), rather than have the user have to guess and provide it.

I am also unsure if it is just an issue with mass matrix initialization, since the step size remains small after the adaptation too. I am still looking into this though, will update this issue if I find something interesting.
It is surprised to me that Stan is fast on this. Given only 1 data, latent `l` will converge to a value with very small variance. Hence `alpha / beta` also converges to a value, which implies that they are highly correlated.

If small step_size is a real problem, then I guess we can scale initial_mass_matrix to let step_size have a higher value at the end of get_reasonable_step_size method. Theoretically, things should be the same but I'll check if it works. Maybe scale step_size will help for precision issue or dual_averaging,...
I just run the program and observe that the small step_size happens at the initial window adaptation, not find_reasonable_step_size. I'll further see how Stan deals with this small step_size problem during adapting phase.
> I'll further see how Stan deals with this small step_size problem during adapting phase.

The scale of the parameters is so vastly different, that I'm not surprised that we need a small step size, but I would have expected that the mass matrix would help neutralize some of that and make it faster. 

My guess is that we need many more slower adaptation windows for mass matrix to be correctly adapted, but in the meantime, we are just stuck with such low step size as to make it impossible to make progress. Stan might have the same issues, but since it has a much faster runtime, it eventually learns a good mass matrix and is not stuck forever. And, PyMC3 probably gets around this problem by constantly adapting the mass matrix.

Resolving patch:
diff --git a/pyro/infer/mcmc/adaptation.py b/pyro/infer/mcmc/adaptation.py
--- a/pyro/infer/mcmc/adaptation.py
+++ b/pyro/infer/mcmc/adaptation.py
@@ -2,6 +2,7 @@
 import torch
 from collections import namedtuple
 
+import pyro
 import pyro.distributions as dist
 from pyro.ops.dual_averaging import DualAveraging
 from pyro.ops.welford import WelfordCovariance
@@ -27,7 +28,7 @@ def __init__(self,
         self.adapt_mass_matrix = adapt_mass_matrix
         self.target_accept_prob = target_accept_prob
         self.is_diag_mass = is_diag_mass
-        self._step_size = step_size if step_size is not None else 1
+        self.step_size = step_size if step_size is not None else 1
         self._adaptation_disabled = not (adapt_step_size or adapt_mass_matrix)
         if adapt_step_size:
             self._step_size_adapt_scheme = DualAveraging()
@@ -85,8 +86,14 @@ def _build_adaptation_schedule(self):
                                                 self._warmup_steps - 1))
         return adaptation_schedule
 
-    def _reset_step_size_adaptation(self, step_size):
-        self._step_size_adapt_scheme.prox_center = math.log(10 * step_size)
+    def reset_step_size_adaptation(self):
+        r"""
+        Finds a reasonable step size and resets step size adaptation scheme.
+        """
+        if self._find_reasonable_step_size is not None:
+            with pyro.validation_enabled(False):
+                self.step_size = self._find_reasonable_step_size()
+        self._step_size_adapt_scheme.prox_center = math.log(10 * self.step_size)
         self._step_size_adapt_scheme.reset()
 
     def _update_step_size(self, accept_prob):
@@ -94,7 +101,7 @@ def _update_step_size(self, accept_prob):
         H = self.target_accept_prob - accept_prob
         self._step_size_adapt_scheme.step(H)
         log_step_size, _ = self._step_size_adapt_scheme.get_state()
-        self._step_size = math.exp(log_step_size)
+        self.step_size = math.exp(log_step_size)
 
     def _update_r_dist(self):
         loc = self._inverse_mass_matrix.new_zeros(self._inverse_mass_matrix.size(0))
@@ -107,19 +114,24 @@ def _update_r_dist(self):
     def _end_adaptation(self):
         if self.adapt_step_size:
             _, log_step_size_avg = self._step_size_adapt_scheme.get_state()
-            self._step_size = math.exp(log_step_size_avg)
+            self.step_size = math.exp(log_step_size_avg)
 
-    def configure(self, warmup_steps, initial_step_size=None, inv_mass_matrix=None):
+    def configure(self, warmup_steps, initial_step_size=None, inv_mass_matrix=None,
+                  find_reasonable_step_size_fn=None):
         r"""
         Model specific properties that are specified when the HMC kernel is setup.
 
         :param warmup_steps: Number of warmup steps that the sampler is initialized with.
-        :param inv_mass_matrix: Initial value of the inverse mass matrix.
         :param initial_step_size: Step size to use to initialize the Dual Averaging scheme.
+        :param inv_mass_matrix: Initial value of the inverse mass matrix.
+        :param find_reasonable_step_size_fn: A callable to find reasonable step size when
+            mass matrix is changed.
         """
         self._warmup_steps = warmup_steps
-        if initial_step_size is not None and self.adapt_step_size:
+        if initial_step_size is not None:
             self.step_size = initial_step_size
+        if find_reasonable_step_size_fn is not None:
+            self._find_reasonable_step_size = find_reasonable_step_size_fn
         if inv_mass_matrix is not None:
             self.inverse_mass_matrix = inv_mass_matrix
         if self.inverse_mass_matrix is None or self.step_size is None:
@@ -144,7 +156,7 @@ def step(self, t, z, accept_prob):
         mass_matrix_adaptation_phase = self.adapt_mass_matrix and \
             (0 < self._current_window < num_windows - 1)
         if self.adapt_step_size:
-            self._update_step_size(accept_prob)
+            self._update_step_size(accept_prob.item())
         if mass_matrix_adaptation_phase:
             z_flat = torch.cat([z[name].reshape(-1) for name in sorted(z)])
             self._mass_matrix_adapt_scheme.update(z_flat.detach())
@@ -158,11 +170,10 @@ def step(self, t, z, accept_prob):
                 self._current_window += 1
                 return
 
-            if self.adapt_step_size:
-                self._reset_step_size_adaptation(self._step_size)
-
             if mass_matrix_adaptation_phase:
                 self.inverse_mass_matrix = self._mass_matrix_adapt_scheme.get_covariance()
+                if self.adapt_step_size:
+                    self.reset_step_size_adaptation()
 
             self._current_window += 1
 
@@ -170,16 +181,6 @@ def step(self, t, z, accept_prob):
     def adaptation_schedule(self):
         return self._adaptation_schedule
 
-    @property
-    def step_size(self):
-        return self._step_size
-
-    @step_size.setter
-    def step_size(self, value):
-        self._step_size = value
-        if self.adapt_step_size:
-            self._reset_step_size_adaptation(value)
-
     @property
     def inverse_mass_matrix(self):
         return self._inverse_mass_matrix
diff --git a/pyro/infer/mcmc/hmc.py b/pyro/infer/mcmc/hmc.py
--- a/pyro/infer/mcmc/hmc.py
+++ b/pyro/infer/mcmc/hmc.py
@@ -104,10 +104,8 @@ def __init__(self,
             self.trajectory_length = step_size * num_steps
         else:
             self.trajectory_length = 2 * math.pi  # from Stan
-        self.adapt_step_size = adapt_step_size
         self._jit_compile = jit_compile
         self._ignore_jit_warnings = ignore_jit_warnings
-        self.full_mass = full_mass
         self._target_accept_prob = 0.8  # from Stan
         # The following parameter is used in find_reasonable_step_size method.
         # In NUTS paper, this threshold is set to a fixed log(0.5).
@@ -142,13 +140,11 @@ def _compute_trace_log_prob(self, model_trace):
         return self._trace_prob_evaluator.log_prob(model_trace)
 
     def _kinetic_energy(self, r):
-        # TODO: revert to `torch.dot` in pytorch==1.0
-        # See: https://github.com/uber/pyro/issues/1458
         r_flat = torch.cat([r[site_name].reshape(-1) for site_name in sorted(r)])
-        if self.full_mass:
-            return 0.5 * (r_flat * (self.inverse_mass_matrix.matmul(r_flat))).sum()
+        if self.inverse_mass_matrix.dim() == 2:
+            return 0.5 * self.inverse_mass_matrix.matmul(r_flat).dot(r_flat)
         else:
-            return 0.5 * (self.inverse_mass_matrix * (r_flat ** 2)).sum()
+            return 0.5 * self.inverse_mass_matrix.dot(r_flat ** 2)
 
     def _potential_energy(self, z):
         if self._jit_compile:
@@ -205,21 +201,23 @@ def _reset(self):
         self._initial_trace = None
         self._has_enumerable_sites = False
         self._trace_prob_evaluator = None
+        self._z_last = None
         self._potential_energy_last = None
         self._z_grads_last = None
         self._warmup_steps = None
 
-    def _find_reasonable_step_size(self, z):
+    def _find_reasonable_step_size(self):
         step_size = self.step_size
 
         # We are going to find a step_size which make accept_prob (Metropolis correction)
         # near the target_accept_prob. If accept_prob:=exp(-delta_energy) is small,
         # then we have to decrease step_size; otherwise, increase step_size.
-        r, _ = self._sample_r(name="r_presample")
-        energy_current = self._energy(z, r)
-        z_new, r_new, z_grads, potential_energy = velocity_verlet(
-            z, r, self._potential_energy, self.inverse_mass_matrix, step_size)
-        energy_new = potential_energy + self._kinetic_energy(r_new)
+        z, potential_energy, z_grads = self._fetch_from_cache()
+        r, _ = self._sample_r(name="r_presample_0")
+        energy_current = self._kinetic_energy(r) + potential_energy
+        z_new, r_new, z_grads_new, potential_energy_new = velocity_verlet(
+            z, r, self._potential_energy, self.inverse_mass_matrix, step_size, z_grads=z_grads)
+        energy_new = self._kinetic_energy(r_new) + potential_energy_new
         delta_energy = energy_new - energy_current
         # direction=1 means keep increasing step_size, otherwise decreasing step_size.
         # Note that the direction is -1 if delta_energy is `NaN` which may be the
@@ -232,11 +230,15 @@ def _find_reasonable_step_size(self, z):
         direction_new = direction
         # keep scale step_size until accept_prob crosses its target
         # TODO: make thresholds for too small step_size or too large step_size
+        t = 0
         while direction_new == direction:
+            t += 1
             step_size = step_size_scale * step_size
-            z_new, r_new, z_grads, potential_energy = velocity_verlet(
-                z, r, self._potential_energy, self.inverse_mass_matrix, step_size)
-            energy_new = potential_energy + self._kinetic_energy(r_new)
+            r, _ = self._sample_r(name="r_presample_{}".format(t))
+            energy_current = self._kinetic_energy(r) + potential_energy
+            z_new, r_new, z_grads_new, potential_energy_new = velocity_verlet(
+                z, r, self._potential_energy, self.inverse_mass_matrix, step_size, z_grads=z_grads)
+            energy_new = self._kinetic_energy(r_new) + potential_energy_new
             delta_energy = energy_new - energy_current
             direction_new = 1 if self._direction_threshold < -delta_energy else -1
         return step_size
@@ -259,18 +261,6 @@ def _guess_max_plate_nesting(self):
                 if frame.vectorized]
         self.max_plate_nesting = -min(dims) if dims else 0
 
-    def _configure_adaptation(self):
-        initial_step_size = None
-        if self.adapt_step_size:
-            z = {name: node["value"].detach() for name, node in self._iter_latent_nodes(self.initial_trace)}
-            for name, transform in self.transforms.items():
-                z[name] = transform(z[name])
-            with pyro.validation_enabled(False):
-                initial_step_size = self._find_reasonable_step_size(z)
-
-        self._adapter.configure(self._warmup_steps,
-                                initial_step_size)
-
     def _sample_r(self, name):
         r_dist = self._adapter.r_dist
         r_flat = pyro.sample(name, r_dist)
@@ -316,6 +306,8 @@ def initial_trace(self):
     @initial_trace.setter
     def initial_trace(self, trace):
         self._initial_trace = trace
+        if self._warmup_steps is not None:  # if setup is already called
+            self._initialize_step_size()
 
     def _initialize_model_properties(self):
         if self.max_plate_nesting is None:
@@ -343,38 +335,48 @@ def _initialize_model_properties(self):
                                                           self._has_enumerable_sites,
                                                           self.max_plate_nesting)
         mass_matrix_size = sum(self._r_numels.values())
-        if self.full_mass:
-            initial_mass_matrix = eye_like(site_value, mass_matrix_size)
-        else:
+        if self._adapter.is_diag_mass:
             initial_mass_matrix = site_value.new_ones(mass_matrix_size)
-        self._adapter.inverse_mass_matrix = initial_mass_matrix
+        else:
+            initial_mass_matrix = eye_like(site_value, mass_matrix_size)
+        self._adapter.configure(self._warmup_steps,
+                                inv_mass_matrix=initial_mass_matrix,
+                                find_reasonable_step_size_fn=self._find_reasonable_step_size)
+        self._initialize_step_size()  # this method also caches z and its potential energy
+
+    def _initialize_step_size(self):
+        z = {name: node["value"].detach()
+             for name, node in self._iter_latent_nodes(self.initial_trace)}
+        # automatically transform `z` to unconstrained space, if needed.
+        for name, transform in self.transforms.items():
+            z[name] = transform(z[name])
+        potential_energy = self._potential_energy(z)
+        self._cache(z, potential_energy, None)
+        if self._adapter.adapt_step_size:
+            self._adapter.reset_step_size_adaptation()
 
     def setup(self, warmup_steps, *args, **kwargs):
         self._warmup_steps = warmup_steps
         self._args = args
         self._kwargs = kwargs
         self._initialize_model_properties()
-        self._configure_adaptation()
 
     def cleanup(self):
         self._reset()
 
-    def _cache(self, potential_energy, z_grads):
+    def _cache(self, z, potential_energy, z_grads):
+        self._z_last = z
         self._potential_energy_last = potential_energy
         self._z_grads_last = z_grads
 
     def _fetch_from_cache(self):
-        return self._potential_energy_last, self._z_grads_last
+        return self._z_last, self._potential_energy_last, self._z_grads_last
 
     def sample(self, trace):
-        z = {name: node["value"].detach() for name, node in self._iter_latent_nodes(trace)}
-        # automatically transform `z` to unconstrained space, if needed.
-        for name, transform in self.transforms.items():
-            z[name] = transform(z[name])
-
+        z, potential_energy, z_grads = self._fetch_from_cache()
         r, _ = self._sample_r(name="r_t={}".format(self._t))
+        energy_current = self._kinetic_energy(r) + potential_energy
 
-        potential_energy, z_grads = self._fetch_from_cache()
         # Temporarily disable distributions args checking as
         # NaNs are expected during step size adaptation
         with optional(pyro.validation_enabled(False), self._t < self._warmup_steps):
@@ -385,8 +387,6 @@ def sample(self, trace):
                                                                               z_grads=z_grads)
             # apply Metropolis correction.
             energy_proposal = self._kinetic_energy(r_new) + potential_energy_new
-            energy_current = self._kinetic_energy(r) + potential_energy if potential_energy is not None \
-                else self._energy(z, r)
         delta_energy = energy_proposal - energy_current
         # Set accept prob to 0.0 if delta_energy is `NaN` which may be
         # the case for a diverging trajectory when using a large step size.
@@ -398,6 +398,7 @@ def sample(self, trace):
         if rand < accept_prob:
             self._accept_cnt += 1
             z = z_new
+            self._cache(z, potential_energy_new, z_grads_new)
 
         if self._t < self._warmup_steps:
             self._adapter.step(self._t, z, accept_prob)
@@ -405,6 +406,7 @@ def sample(self, trace):
         self._t += 1
 
         # get trace with the constrained values for `z`.
+        z = z.copy()
         for name, transform in self.transforms.items():
             z[name] = transform.inv(z[name])
         return self._get_trace(z)
diff --git a/pyro/infer/mcmc/nuts.py b/pyro/infer/mcmc/nuts.py
--- a/pyro/infer/mcmc/nuts.py
+++ b/pyro/infer/mcmc/nuts.py
@@ -134,16 +134,13 @@ def _is_turning(self, r_left, r_right, r_sum):
         # We follow the strategy in Section A.4.2 of [2] for this implementation.
         r_left_flat = torch.cat([r_left[site_name].reshape(-1) for site_name in sorted(r_left)])
         r_right_flat = torch.cat([r_right[site_name].reshape(-1) for site_name in sorted(r_right)])
-        # TODO: change to torch.dot for pytorch 1.0
-        if self.full_mass:
-            if (((r_sum - r_left_flat) * (self.inverse_mass_matrix.matmul(r_left_flat)))
-                    .sum() > 0 and
-                    ((r_sum - r_right_flat) * (self.inverse_mass_matrix.matmul(r_right_flat)))
-                    .sum() > 0):
+        if self.inverse_mass_matrix.dim() == 2:
+            if (self.inverse_mass_matrix.matmul(r_left_flat).dot(r_sum - r_left_flat) > 0 and
+                    self.inverse_mass_matrix.matmul(r_right_flat).dot(r_sum - r_right_flat) > 0):
                 return False
         else:
-            if ((self.inverse_mass_matrix * (r_sum - r_left_flat) * r_left_flat).sum() > 0 and
-                    (self.inverse_mass_matrix * (r_sum - r_right_flat) * r_right_flat).sum() > 0):
+            if (self.inverse_mass_matrix.mul(r_left_flat).dot(r_sum - r_left_flat) > 0 and
+                    self.inverse_mass_matrix.mul(r_right_flat).dot(r_sum - r_right_flat) > 0):
                 return False
         return True
 
@@ -256,14 +253,9 @@ def _build_tree(self, z, r, z_grads, log_slice, direction, tree_depth, energy_cu
                          sum_accept_probs, num_proposals)
 
     def sample(self, trace):
-        z = {name: node["value"].detach() for name, node in self._iter_latent_nodes(trace)}
-        potential_energy, z_grads = self._fetch_from_cache()
-        # automatically transform `z` to unconstrained space, if needed.
-        for name, transform in self.transforms.items():
-            z[name] = transform(z[name])
+        z, potential_energy, z_grads = self._fetch_from_cache()
         r, r_flat = self._sample_r(name="r_t={}".format(self._t))
-        energy_current = self._kinetic_energy(r) + potential_energy if potential_energy is not None \
-            else self._energy(z, r)
+        energy_current = self._kinetic_energy(r) + potential_energy
 
         # Ideally, following a symplectic integrator trajectory, the energy is constant.
         # In that case, we can sample the proposal uniformly, and there is no need to use "slice".
@@ -295,6 +287,8 @@ def sample(self, trace):
         z_left_grads = z_right_grads = z_grads
         accepted = False
         r_sum = r_flat
+        sum_accept_probs = 0.
+        num_proposals = 0
         if self.use_multinomial_sampling:
             tree_weight = energy_current.new_zeros(())
         else:
@@ -304,7 +298,8 @@ def sample(self, trace):
         # NaNs are expected during step size adaptation.
         with optional(pyro.validation_enabled(False), self._t < self._warmup_steps):
             # doubling process, stop when turning or diverging
-            for tree_depth in range(self._max_tree_depth + 1):
+            tree_depth = 0
+            while tree_depth < self._max_tree_depth:
                 direction = pyro.sample("direction_t={}_treedepth={}".format(self._t, tree_depth),
                                         dist.Bernoulli(probs=torch.ones(1) * 0.5))
                 direction = int(direction.item())
@@ -322,9 +317,14 @@ def sample(self, trace):
                     r_left = new_tree.r_left
                     z_left_grads = new_tree.z_left_grads
 
+                sum_accept_probs = sum_accept_probs + new_tree.sum_accept_probs
+                num_proposals = num_proposals + new_tree.num_proposals
+
                 if new_tree.turning or new_tree.diverging:  # stop doubling
                     break
 
+                tree_depth += 1
+
                 if self.use_multinomial_sampling:
                     new_tree_prob = (new_tree.weight - tree_weight).exp()
                 else:
@@ -335,7 +335,7 @@ def sample(self, trace):
                 if rand < new_tree_prob:
                     accepted = True
                     z = new_tree.z_proposal
-                    self._cache(new_tree.z_proposal_pe, new_tree.z_proposal_grads)
+                    self._cache(z, new_tree.z_proposal_pe, new_tree.z_proposal_grads)
 
                 r_sum = r_sum + new_tree.r_sum
                 if self._is_turning(r_left, r_right, r_sum):  # stop doubling
@@ -347,7 +347,7 @@ def sample(self, trace):
                         tree_weight = tree_weight + new_tree.weight
 
         if self._t < self._warmup_steps:
-            accept_prob = new_tree.sum_accept_probs / new_tree.num_proposals
+            accept_prob = sum_accept_probs / num_proposals
             self._adapter.step(self._t, z, accept_prob)
 
         if accepted:
@@ -355,6 +355,7 @@ def sample(self, trace):
 
         self._t += 1
         # get trace with the constrained values for `z`.
+        z = z.copy()
         for name, transform in self.transforms.items():
             z[name] = transform.inv(z[name])
         return self._get_trace(z)


Test patch:
diff --git a/tests/infer/mcmc/test_hmc.py b/tests/infer/mcmc/test_hmc.py
--- a/tests/infer/mcmc/test_hmc.py
+++ b/tests/infer/mcmc/test_hmc.py
@@ -272,7 +272,7 @@ def test_initial_trace(monkeypatch):
     true_coefs = torch.arange(1., dim + 1.)
     labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample()
     # mock values to replay - from right to left
-    trace_log_prob_replay = [-5, float('NaN'), float('Inf')]
+    trace_log_prob_replay = [-5, -5, float('NaN'), float('Inf')]
 
     def model(data):
         coefs_mean = torch.zeros(dim)