Divergence checking does not maintain time-reversibility

[howsiyu]

As per title. This can cause sampling bias when we have jump near max_delta_energy. Here's a concrete example translated from @nhuurre 's stan code at https://discourse.mc-stan.org/t/divergence-check-does-not-satisfy-time-reversibility/33738.

import jax
import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS
import seaborn as sns

# Sampling is biased when jump is near `max_delta_energy`.
jump = 1000.0

def model():
    x = numpyro.sample("x", dist.Normal())
    numpyro.factor("jump", jnp.where(x > 0, jump, 0))

nuts_kernel = NUTS(model, target_accept_prob=0.6)
mcmc = MCMC(nuts_kernel, num_warmup=2_000, num_samples=20_000)
mcmc.run(jax.random.key(0))

sns.histplot(x=mcmc.get_samples()['x'], stat="density")
xs = jnp.linspace(0.0, 3.5, 100)
sns.lineplot(x=xs, y=jnp.exp(jax.vmap(dist.HalfNormal().log_prob)(xs)))

[fehiepsi]

I guess when a divergence happens, we will ignore the trajectory and diagnose the model to find the issue. So I don't think we need to worry about the time-reversibility here.

[howsiyu]

when a divergence happens, we will ignore the trajectory and diagnose the model to find the issue.

If that's the case, shouldn't numpyro stop immedately once it encounters a divergence post-warmup?

[martinjankowiak]

maybe numpyro should stop immediately because floating point errors violate time-reversibility ; )