Run variational inference¶

Through our PyMC interface, we also gain access to approximate posteriors via Variational Inference. This guide shows how to run variational inference through HSSM, and when to prefer it over MCMC.

This guide covers four things: the basics of running VI through HSSM (.vi(), its outputs, convergence checks), when to prefer VI over MCMC, a hierarchical case study comparing both under centered and non-centered parameterizations, and an appendix driving PyMC's VI machinery directly.

In [1]:

Copied!





import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm

import hssm
import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm

import hssm

Load Data and Specify model¶

In [2]:

Copied!

cav_data = hssm.load_data("cavanagh_theta")
cav_model = hssm.HSSM(data=cav_data, model="angle")
cav_data = hssm.load_data("cavanagh_theta")
cav_model = hssm.HSSM(data=cav_data, model="angle")

Model initialized successfully.

Inference¶

We will run MCMC and VI here to contrast results.

Run MCMC¶

In [3]:

Copied!





# Slight adjustment to initial values
initvals_tmp = cav_model.initvals
initvals_tmp["theta"] = 0.1


mcmc_idata = cav_model.sample(
    chains=2, tune=500, draws=500, sampler="numpyro", initvals=initvals_tmp
)
# Slight adjustment to initial values
initvals_tmp = cav_model.initvals
initvals_tmp["theta"] = 0.1


mcmc_idata = cav_model.sample(
    chains=2, tune=500, draws=500, sampler="numpyro", initvals=initvals_tmp
)

bambi/backend/pymc.py:232: UserWarning: `init='adapt_diag'` is ignored by `nuts_sampler='numpyro'`; the external sampler uses its own initialization.
  idata = pm.sample(
NUTS[numpyro]: [t, theta, z, a, v]

pymc/sampling/jax.py:463: UserWarning: There are not enough devices to run parallel chains: expected 2 but got 1. Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(2)` at the beginning of your program. You can double-check how many devices are available in your system using `jax.local_device_count()`.
  pmap_numpyro = MCMC(

  0%|          | 0/1000 [00:00<?, ?it/s]

warmup:   5%|▌         | 50/1000 [00:01<00:34, 27.91it/s, 7 steps of size 9.36e-03. acc. prob=0.74]

warmup:  10%|█         | 100/1000 [00:02<00:24, 36.19it/s, 7 steps of size 4.60e-03. acc. prob=0.76]

warmup:  15%|█▌        | 150/1000 [00:06<00:42, 19.79it/s, 3 steps of size 1.54e-01. acc. prob=0.78]

warmup:  20%|██        | 200/1000 [00:07<00:31, 25.39it/s, 3 steps of size 1.09e-01. acc. prob=0.78]

warmup:  25%|██▌       | 250/1000 [00:08<00:24, 30.14it/s, 15 steps of size 2.33e-01. acc. prob=0.78]

warmup:  30%|███       | 300/1000 [00:10<00:24, 28.95it/s, 7 steps of size 4.01e-01. acc. prob=0.78]

warmup:  35%|███▌      | 350/1000 [00:12<00:21, 30.69it/s, 7 steps of size 4.03e-01. acc. prob=0.78]

warmup:  40%|████      | 400/1000 [00:13<00:18, 32.05it/s, 7 steps of size 2.11e-01. acc. prob=0.78]

warmup:  45%|████▌     | 450/1000 [00:14<00:16, 34.20it/s, 15 steps of size 4.52e-01. acc. prob=0.79]

warmup:  50%|█████     | 500/1000 [00:16<00:15, 33.27it/s, 7 steps of size 2.67e-01. acc. prob=0.79]

sample:  55%|█████▌    | 550/1000 [00:17<00:13, 34.35it/s, 7 steps of size 2.67e-01. acc. prob=0.93]

sample:  60%|██████    | 600/1000 [00:19<00:11, 33.83it/s, 15 steps of size 2.67e-01. acc. prob=0.92]

sample:  65%|██████▌   | 650/1000 [00:20<00:10, 34.43it/s, 3 steps of size 2.67e-01. acc. prob=0.92]

sample:  70%|███████   | 700/1000 [00:22<00:08, 34.68it/s, 31 steps of size 2.67e-01. acc. prob=0.92]

sample:  75%|███████▌  | 750/1000 [00:23<00:07, 33.71it/s, 7 steps of size 2.67e-01. acc. prob=0.92]

sample:  80%|████████  | 800/1000 [00:25<00:05, 33.75it/s, 7 steps of size 2.67e-01. acc. prob=0.92]

sample:  85%|████████▌ | 850/1000 [00:26<00:04, 35.65it/s, 11 steps of size 2.67e-01. acc. prob=0.92]

sample:  90%|█████████ | 900/1000 [00:28<00:02, 33.85it/s, 7 steps of size 2.67e-01. acc. prob=0.92]

sample:  95%|█████████▌| 950/1000 [00:29<00:01, 33.51it/s, 31 steps of size 2.67e-01. acc. prob=0.92]

sample: 100%|██████████| 1000/1000 [00:31<00:00, 32.92it/s, 13 steps of size 2.67e-01. acc. prob=0.92]

sample: 100%|██████████| 1000/1000 [00:31<00:00, 31.99it/s, 13 steps of size 2.67e-01. acc. prob=0.92]

  0%|          | 0/1000 [00:00<?, ?it/s]

warmup:   5%|▌         | 50/1000 [00:01<00:24, 39.41it/s, 7 steps of size 6.86e-03. acc. prob=0.74]

warmup:  10%|█         | 100/1000 [00:02<00:22, 40.88it/s, 15 steps of size 1.30e-02. acc. prob=0.76]

warmup:  15%|█▌        | 150/1000 [00:04<00:29, 28.67it/s, 31 steps of size 1.62e-01. acc. prob=0.78]

warmup:  20%|██        | 200/1000 [00:06<00:27, 28.73it/s, 7 steps of size 4.73e-01. acc. prob=0.78]

warmup:  25%|██▌       | 250/1000 [00:07<00:23, 32.26it/s, 7 steps of size 1.49e-01. acc. prob=0.78]

warmup:  30%|███       | 300/1000 [00:09<00:22, 30.67it/s, 15 steps of size 5.08e-01. acc. prob=0.78]

warmup:  35%|███▌      | 350/1000 [00:10<00:20, 32.00it/s, 15 steps of size 6.45e-01. acc. prob=0.78]

warmup:  40%|████      | 400/1000 [00:12<00:16, 35.83it/s, 11 steps of size 2.68e-01. acc. prob=0.79]

warmup:  45%|████▌     | 450/1000 [00:13<00:14, 37.63it/s, 3 steps of size 2.20e-01. acc. prob=0.79]

warmup:  50%|█████     | 500/1000 [00:15<00:14, 33.71it/s, 3 steps of size 2.54e-01. acc. prob=0.78]

sample:  55%|█████▌    | 550/1000 [00:16<00:13, 33.45it/s, 15 steps of size 2.54e-01. acc. prob=0.94]

sample:  60%|██████    | 600/1000 [00:18<00:11, 33.42it/s, 31 steps of size 2.54e-01. acc. prob=0.94]

sample:  65%|██████▌   | 650/1000 [00:19<00:10, 33.57it/s, 15 steps of size 2.54e-01. acc. prob=0.94]

sample:  70%|███████   | 700/1000 [00:21<00:09, 31.67it/s, 23 steps of size 2.54e-01. acc. prob=0.94]

sample:  75%|███████▌  | 750/1000 [00:22<00:07, 32.01it/s, 3 steps of size 2.54e-01. acc. prob=0.94]

sample:  80%|████████  | 800/1000 [00:24<00:06, 31.07it/s, 19 steps of size 2.54e-01. acc. prob=0.94]

sample:  85%|████████▌ | 850/1000 [00:26<00:04, 31.88it/s, 11 steps of size 2.54e-01. acc. prob=0.94]

sample:  90%|█████████ | 900/1000 [00:27<00:03, 31.74it/s, 7 steps of size 2.54e-01. acc. prob=0.94]

sample:  95%|█████████▌| 950/1000 [00:29<00:01, 32.82it/s, 15 steps of size 2.54e-01. acc. prob=0.94]

sample: 100%|██████████| 1000/1000 [00:30<00:00, 32.68it/s, 15 steps of size 2.54e-01. acc. prob=0.94]

sample: 100%|██████████| 1000/1000 [00:30<00:00, 32.72it/s, 15 steps of size 2.54e-01. acc. prob=0.94]

We recommend running at least 4 chains for robust computation of convergence diagnostics

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

In [4]:

Copied!

az.summary(mcmc_idata)
az.summary(mcmc_idata)

Out[4]:

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
z	0.5049	0.006	0.5	0.51	502	618	1.01	0.00027	0.00019
t	0.2849	0.0094	0.27	0.3	377	539	1.00	0.00048	0.00035
a	1.303	0.022	1.3	1.3	343	407	1.00	0.0012	0.00081
v	0.361	0.021	0.33	0.39	543	693	1.00	0.00091	0.00061
theta	0.226	0.0129	0.21	0.25	364	431	1.00	0.00068	0.00047

Run VI¶

In [5]:

Copied!

vi_idata = cav_model.vi(niter=50000, method="fullrank_advi")
vi_idata = cav_model.vi(niter=50000, method="fullrank_advi")

Using MCMC starting point defaults.

pytensor/link/jax/linker.py:50: UserWarning: The RandomType SharedVariables [RNG(<Generator(PCG64) at 0x138A66880>)] will not be used in the compiled JAX graph. Instead a copy will be used.
  warnings.warn(

Finished [100%]: Average Loss = 6,071.1

Inspect Outputs¶

From our variational inference runs, we extract two objects.

An xarray.DataTree object stored under cav_model.vi_idata. This stores a slightly cleaned up posterior sample, constructed by sampling from the variational posterior.
A PyMC variational Approximation object stored under cav_model.vi_approx that holds the variational posterior object itself. This is a rich structure and it is beyond the purpose of this tutorial to illustrate all it's details. Amongst other things you will be able to inspect the loss history and take samples such as those stored under cav_model.vi_idata.

`.vi_idata`¶

The approximate variational posterior as an xarray.DataTree.

`.vi_approx`¶

The approximate variational posterior as a PyMC variational Approximation object.

We can take draws from the posterior with the .sample() method.

In [7]:

Copied!

with cav_model.pymc_model:
    cav_model.vi_approx.sample(draws=1000)
with cav_model.pymc_model:
    cav_model.vi_approx.sample(draws=1000)

The .hist attribute stores the loss history. We can plot this to see how the loss function converged.

In [8]:

Copied!

plt.plot(cav_model.vi_approx.hist)
plt.xlabel("Iteration")
plt.ylabel("Loss")
plt.plot(cav_model.vi_approx.hist)
plt.xlabel("Iteration")
plt.ylabel("Loss")

Out[8]:

Text(0, 0.5, 'Loss')

No description has been provided for this image

Contrast outputs between MCMC and VI¶

In [9]:

Copied!





plot_matrix = az.plot_pair(
    cav_model.traces,
    visuals={
        "scatter": dict(alpha=0.01, color="blue"),
        "dist": dict(alpha=0.6, color="blue"),
    },
)
az.plot_pair(
    cav_model.vi_idata,
    plot_matrix=plot_matrix,
    visuals={
        "scatter": dict(alpha=0.04, color="red"),
        "dist": dict(alpha=0.8, color="red"),
    },
)
plot_matrix = az.plot_pair(
    cav_model.traces,
    visuals={
        "scatter": dict(alpha=0.01, color="blue"),
        "dist": dict(alpha=0.6, color="blue"),
    },
)
az.plot_pair(
    cav_model.vi_idata,
    plot_matrix=plot_matrix,
    visuals={
        "scatter": dict(alpha=0.04, color="red"),
        "dist": dict(alpha=0.8, color="red"),
    },
)

Out[9]:

<arviz_plots.plot_matrix.PlotMatrix at 0x13ae61e80>

When to use VI — and when to stick with MCMC¶

Variational inference replaces sampling with deterministic optimization: it fits the parameters of an approximating distribution by maximizing the ELBO. That buys speed and scalability, and costs exactness. A practical decision guide:

Reach for VI when…

you are iterating on a large (e.g. hierarchical) model and MCMC turnaround is the bottleneck — VI is excellent for prototyping model structure before a final MCMC run;
you need point estimates and rough uncertainty on a model too large to sample comfortably;
you can verify the approximation: compare a VI fit against a short MCMC run at least once (as we do below).

Stick with MCMC when…

the uncertainty itself is the scientific quantity (group-level sigmas, individual differences, Bayes-factor-adjacent decisions) — mean-field VI systematically underestimates posterior variance and drops parameter correlations; full-rank helps with correlations but is costlier and still approximate;
posteriors are likely multimodal or heavy-tailed;
you are producing final, publication-grade inferences.

Always, with VI:

check the loss history (.vi_approx.hist) for convergence — the required iteration count is hard to guess in advance;
remember that a converged loss means the optimizer finished, not that the approximation is good — the case study below shows one setting where VI matches MCMC closely and one where it is visibly overconfident;
prefer fullrank_advi over advi when parameter correlations matter and the parameter count allows it.

Case study: VI on a hierarchical DDM¶

We now put VI to work where it matters most — a hierarchical model, where MCMC is slowest. We simulate 20 participants × 120 trials of DDM data, then fit the same model twice (centered and non-centered parameterization), each with both VI and MCMC, and compare parameter recovery against the known ground truth.

Thanks to Guillaume Pagnier, PhD, for the initial version of this case study.

Utilities¶

In [11]:

Copied!





def process_idata_for_plotting(
    idata: az.InferenceData, parameter_matrix: pd.DataFrame, model: str
) -> pd.DataFrame:
    """Process inference data and parameter matrix into a dataframe for plotting.

    Parameters
    ----------
    idata : az.InferenceData
        Inference data containing posterior samples
    parameter_matrix : pd.DataFrame
        DataFrame containing true parameter values
    model : str
        Name of model to get parameters for from defaults

    Returns
    -------
    pd.DataFrame
        DataFrame containing processed posterior means, HDIs and true parameters
        ready for plotting
    """
    # Get Posterior Means VI
    params_df_mean = pd.DataFrame(
        {
            param: idata.posterior[param].mean(dim=["chain", "draw"]).values
            for param in hssm.defaults.default_model_config[model]["list_params"]
        }
    )
    params_df_mean.columns = [f"{param}_mean" for param in params_df_mean.columns]

    # Get Posterior HDIs VI
    params_df_hdi = (
        az.hdi(
            idata.posterior.ds[
                hssm.defaults.default_model_config[model]["list_params"]
            ],
            prob=0.95,
        )
        .to_dataframe()
        .reset_index()
        .pivot(
            index="__obs__",
            columns="ci_bound",
            values=hssm.defaults.default_model_config[model]["list_params"],
        )
    )

    # Get rid of multiindex
    params_df_hdi.columns = ["_".join(col + ("hdi",)) for col in params_df_hdi.columns]
    params_df_hdi = params_df_hdi.reset_index(drop=True)

    # Combine data
    data_processed = pd.concat([simDataDDM, params_df_mean, params_df_hdi], axis=1)

    # Make plotting ready
    plot_df = pd.concat(
        [
            data_processed.drop(columns=["rt", "response"])
            .drop_duplicates()
            .reset_index(drop=True),
            parameter_matrix.drop(columns=["participant_id"]),
        ],
        axis=1,
    )
    return plot_df
def process_idata_for_plotting(
    idata: az.InferenceData, parameter_matrix: pd.DataFrame, model: str
) -> pd.DataFrame:
    """Process inference data and parameter matrix into a dataframe for plotting.

    Parameters
    ----------
    idata : az.InferenceData
        Inference data containing posterior samples
    parameter_matrix : pd.DataFrame
        DataFrame containing true parameter values
    model : str
        Name of model to get parameters for from defaults

    Returns
    -------
    pd.DataFrame
        DataFrame containing processed posterior means, HDIs and true parameters
        ready for plotting
    """
    # Get Posterior Means VI
    params_df_mean = pd.DataFrame(
        {
            param: idata.posterior[param].mean(dim=["chain", "draw"]).values
            for param in hssm.defaults.default_model_config[model]["list_params"]
        }
    )
    params_df_mean.columns = [f"{param}_mean" for param in params_df_mean.columns]

    # Get Posterior HDIs VI
    params_df_hdi = (
        az.hdi(
            idata.posterior.ds[
                hssm.defaults.default_model_config[model]["list_params"]
            ],
            prob=0.95,
        )
        .to_dataframe()
        .reset_index()
        .pivot(
            index="__obs__",
            columns="ci_bound",
            values=hssm.defaults.default_model_config[model]["list_params"],
        )
    )

    # Get rid of multiindex
    params_df_hdi.columns = ["_".join(col + ("hdi",)) for col in params_df_hdi.columns]
    params_df_hdi = params_df_hdi.reset_index(drop=True)

    # Combine data
    data_processed = pd.concat([simDataDDM, params_df_mean, params_df_hdi], axis=1)

    # Make plotting ready
    plot_df = pd.concat(
        [
            data_processed.drop(columns=["rt", "response"])
            .drop_duplicates()
            .reset_index(drop=True),
            parameter_matrix.drop(columns=["participant_id"]),
        ],
        axis=1,
    )
    return plot_df

Simulate Dataset¶

We simulate a dataset of 20 participants, each with 120 trials.

In [12]:

Copied!





# Group level parameters
v_mu = 0.2
v_sigma = 0.3
a_mu = 1
a_sigma = 0.2
z_mu = 0.5
z_sigma = 0.01
t_mu = 0.4
t_sigma = 0.00

# Make hierarchical dataset
parameter_matrix = pd.DataFrame(
    {
        "participant_id": [f"subj_{str(i).zfill(2)}" for i in range(1, 21)],
        "v_true": np.sort(np.random.normal(loc=v_mu, scale=v_sigma, size=20)).round(1),
        "a_true": np.sort(np.random.normal(loc=a_mu, scale=a_sigma, size=20)).round(1),
        "z_true": np.sort(np.random.normal(loc=z_mu, scale=z_sigma, size=20)).round(1),
        "t_true": np.sort(np.random.normal(loc=t_mu, scale=t_sigma, size=20)).round(1),
        "nTrials": [120] * 20,
    }
)


dfs = []
for _, row in parameter_matrix.iterrows():
    df = hssm.simulate_data(
        model="ddm",
        theta=dict(v=row["v_true"], a=row["a_true"], z=row["z_true"], t=row["t_true"]),
        size=row["nTrials"],
    )
    df["participant_id"] = row["participant_id"]
    dfs.append(df)
    simDataDDM = pd.concat(dfs, ignore_index=True)
# Group level parameters
v_mu = 0.2
v_sigma = 0.3
a_mu = 1
a_sigma = 0.2
z_mu = 0.5
z_sigma = 0.01
t_mu = 0.4
t_sigma = 0.00

# Make hierarchical dataset
parameter_matrix = pd.DataFrame(
    {
        "participant_id": [f"subj_{str(i).zfill(2)}" for i in range(1, 21)],
        "v_true": np.sort(np.random.normal(loc=v_mu, scale=v_sigma, size=20)).round(1),
        "a_true": np.sort(np.random.normal(loc=a_mu, scale=a_sigma, size=20)).round(1),
        "z_true": np.sort(np.random.normal(loc=z_mu, scale=z_sigma, size=20)).round(1),
        "t_true": np.sort(np.random.normal(loc=t_mu, scale=t_sigma, size=20)).round(1),
        "nTrials": [120] * 20,
    }
)


dfs = []
for _, row in parameter_matrix.iterrows():
    df = hssm.simulate_data(
        model="ddm",
        theta=dict(v=row["v_true"], a=row["a_true"], z=row["z_true"], t=row["t_true"]),
        size=row["nTrials"],
    )
    df["participant_id"] = row["participant_id"]
    dfs.append(df)
    simDataDDM = pd.concat(dfs, ignore_index=True)

Hierarchical Model: Centered Parameterization¶

We first instantiate the generative model using HSSM syntax. We will fit this model to the synthetic data using two approaches 1. Variational inference 2. MCMC

In [13]:

Copied!





# Generative model
mSimCentered = hssm.HSSM(
    data=simDataDDM,
    p_outlier=0.01,
    prior_settings="safe",
    noncentered=False,
    model="ddm",
    loglik_kind="approx_differentiable",
    include=[
        {
            "name": "v",
            "formula": "v ~ 0 +(1|participant_id)",
            "prior": {
                "1|participant_id": {
                    "name": "Normal",
                    "mu": {
                        "name": "Normal",
                        "mu": 1.3,
                        "sigma": 0.3,
                    },
                    "sigma": {"name": "HalfNormal", "sigma": 0.2},
                },
            },
        },
        {
            "name": "a",
            "formula": "a ~ 0 + (1|participant_id)",
            "prior": {
                "1|participant_id": {
                    "name": "Normal",
                    "mu": {"name": "Gamma", "mu": 1, "sigma": 0.2},
                    "sigma": {"name": "HalfNormal", "sigma": 0.2},
                },
            },
        },
        {
            "name": "z",
            "formula": "z ~ 0 + (1|participant_id)",
            "prior": {
                "1|participant_id": {
                    "name": "Normal",
                    "mu": {"name": "Beta", "alpha": 10, "beta": 10},
                    "sigma": {"name": "HalfNormal", "sigma": 0.01},
                },
            },
        },
        {
            "name": "t",
            "formula": "t ~ 0 + (1|participant_id)",
            "prior": {
                "1|participant_id": {
                    "name": "Normal",
                    "mu": {
                        "name": "Normal",
                        "mu": 0.4,
                        "sigma": 0.1,
                    },
                    "sigma": {"name": "HalfNormal", "sigma": 0.01},
                },
            },
        },
    ],
)
# Generative model
mSimCentered = hssm.HSSM(
    data=simDataDDM,
    p_outlier=0.01,
    prior_settings="safe",
    noncentered=False,
    model="ddm",
    loglik_kind="approx_differentiable",
    include=[
        {
            "name": "v",
            "formula": "v ~ 0 +(1|participant_id)",
            "prior": {
                "1|participant_id": {
                    "name": "Normal",
                    "mu": {
                        "name": "Normal",
                        "mu": 1.3,
                        "sigma": 0.3,
                    },
                    "sigma": {"name": "HalfNormal", "sigma": 0.2},
                },
            },
        },
        {
            "name": "a",
            "formula": "a ~ 0 + (1|participant_id)",
            "prior": {
                "1|participant_id": {
                    "name": "Normal",
                    "mu": {"name": "Gamma", "mu": 1, "sigma": 0.2},
                    "sigma": {"name": "HalfNormal", "sigma": 0.2},
                },
            },
        },
        {
            "name": "z",
            "formula": "z ~ 0 + (1|participant_id)",
            "prior": {
                "1|participant_id": {
                    "name": "Normal",
                    "mu": {"name": "Beta", "alpha": 10, "beta": 10},
                    "sigma": {"name": "HalfNormal", "sigma": 0.01},
                },
            },
        },
        {
            "name": "t",
            "formula": "t ~ 0 + (1|participant_id)",
            "prior": {
                "1|participant_id": {
                    "name": "Normal",
                    "mu": {
                        "name": "Normal",
                        "mu": 0.4,
                        "sigma": 0.1,
                    },
                    "sigma": {"name": "HalfNormal", "sigma": 0.01},
                },
            },
        },
    ],
)

Model initialized successfully.

In [14]:

Copied!

mSimCentered.graph()
mSimCentered.graph()

Out[14]:

Fit VI¶

We fit with .vi() exactly as in the basics section above — the one new ingredient is a tuned optimizer (pm.adamax). See the decision guide for what VI trades away for this speed.

In [15]:

Copied!





# VI
vi_idata = mSimCentered.vi(
    niter=20000,
    method="advi",
    obj_optimizer=pm.adamax(learning_rate=0.01),
)
with mSimCentered.pymc_model:
    mSimCenteredVIObject = mSimCentered.vi_approx.sample(draws=1000)
# VI
vi_idata = mSimCentered.vi(
    niter=20000,
    method="advi",
    obj_optimizer=pm.adamax(learning_rate=0.01),
)
with mSimCentered.pymc_model:
    mSimCenteredVIObject = mSimCentered.vi_approx.sample(draws=1000)

Using MCMC starting point defaults.

pytensor/link/jax/linker.py:50: UserWarning: The RandomType SharedVariables [RNG(<Generator(PCG64) at 0x13A1DDA80>)] will not be used in the compiled JAX graph. Instead a copy will be used.
  warnings.warn(

Finished [100%]: Average Loss = 3,316.3

Before looking at the posteriors we must ensure that the model successfully "converged".

In [16]:

Copied!





# Loss plot
plt.plot(mSimCentered.vi_approx.hist)
plt.xlabel("Iteration")
plt.ylabel("Loss")
plt.title("VI iteration loss")
plt.show()
# Loss plot
plt.plot(mSimCentered.vi_approx.hist)
plt.xlabel("Iteration")
plt.ylabel("Loss")
plt.title("VI iteration loss")
plt.show()

Once we are satisfied the loss is acceptable, we can take a look at the mean of the group parameters. These should match our generative parameters above.

In [17]:

Copied!





summary_vi = az.summary(
    mSimCenteredVIObject.posterior,
    var_names=[r".*_mu", r".*_sigma"],
    filter_vars="regex",
).sort_index()
summary_vi
summary_vi = az.summary(
    mSimCenteredVIObject.posterior,
    var_names=[r".*_mu", r".*_sigma"],
    filter_vars="regex",
).sort_index()
summary_vi

Out[17]:

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
a_1\|participant_id_mu	0.954	0.044	0.89	1	751	999	nan	0.0016	0.0011
a_1\|participant_id_sigma	0.217	0.033	0.17	0.27	878	884	nan	0.0011	0.00088
t_1\|participant_id_mu	0.4071	0.0026	0.4	0.41	1021	853	nan	8.3e-05	5.9e-05
t_1\|participant_id_sigma	0.0127	0.00193	0.0097	0.016	1112	1072	nan	5.8e-05	4.3e-05
v_1\|participant_id_mu	0.286	0.066	0.18	0.39	894	877	nan	0.0022	0.0016
v_1\|participant_id_sigma	0.309	0.048	0.24	0.39	944	915	nan	0.0016	0.0012
z_1\|participant_id_mu	0.5068	0.0033	0.5	0.51	818	876	nan	0.00011	8.1e-05
z_1\|participant_id_sigma	0.0145	0.0022	0.011	0.018	871	916	nan	7.6e-05	5.7e-05

Fit MCMC¶

Now the same model with MCMC — slower, but the reference for what the posterior actually looks like.

In [18]:

Copied!





# MCMC
mSimCenteredSampled = mSimCentered.sample(
    sampler="numpyro", cores=4, chains=4, draws=250, tune=500, mp_ctx="forkserver"
)

mSimCentered.sample_posterior_predictive(dt=mSimCenteredSampled)
# MCMC
mSimCenteredSampled = mSimCentered.sample(
    sampler="numpyro", cores=4, chains=4, draws=250, tune=500, mp_ctx="forkserver"
)

mSimCentered.sample_posterior_predictive(dt=mSimCenteredSampled)

Using default initvals.

bambi/backend/pymc.py:232: UserWarning: `init='adapt_diag'` is ignored by `nuts_sampler='numpyro'`; the external sampler uses its own initialization.
  idata = pm.sample(
NUTS[numpyro]: [v_1|participant_id_mu, v_1|participant_id_sigma, v_1|participant_id, a_1|participant_id_mu, a_1|participant_id_sigma, a_1|participant_id, z_1|participant_id_mu, z_1|participant_id_sigma, z_1|participant_id, t_1|participant_id_mu, t_1|participant_id_sigma, t_1|participant_id]

pymc/sampling/jax.py:463: UserWarning: There are not enough devices to run parallel chains: expected 4 but got 1. Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(4)` at the beginning of your program. You can double-check how many devices are available in your system using `jax.local_device_count()`.
  pmap_numpyro = MCMC(

  0%|          | 0/750 [00:00<?, ?it/s]

warmup:   0%|          | 0/750 [00:17<?, ?it/s, 127 steps of size 4.66e-03. acc. prob=0.71]

warmup:   5%|▍         | 35/750 [00:17<05:55,  2.01it/s, 223 steps of size 2.38e-03. acc. prob=0.71]

warmup:   5%|▍         | 36/750 [00:17<05:53,  2.02it/s, 255 steps of size 4.10e-03. acc. prob=0.72]

warmup:   5%|▌         | 38/750 [00:18<05:53,  2.02it/s, 511 steps of size 2.28e-03. acc. prob=0.71]

warmup:   5%|▌         | 39/750 [00:19<05:50,  2.03it/s, 255 steps of size 3.98e-03. acc. prob=0.72]

warmup:   5%|▌         | 40/750 [00:19<05:33,  2.13it/s, 127 steps of size 4.92e-03. acc. prob=0.72]

warmup:   6%|▌         | 42/750 [00:20<05:33,  2.12it/s, 511 steps of size 1.19e-03. acc. prob=0.71]

warmup:   6%|▌         | 43/750 [00:21<06:12,  1.90it/s, 511 steps of size 1.76e-03. acc. prob=0.72]

warmup:   6%|▌         | 44/750 [00:21<06:04,  1.93it/s, 255 steps of size 2.96e-03. acc. prob=0.72]

warmup:   6%|▌         | 45/750 [00:22<05:28,  2.15it/s, 127 steps of size 2.85e-03. acc. prob=0.73]

warmup:   6%|▌         | 46/750 [00:22<04:54,  2.39it/s, 127 steps of size 2.40e-03. acc. prob=0.73]

warmup:   6%|▋         | 47/750 [00:22<04:25,  2.65it/s, 127 steps of size 4.09e-03. acc. prob=0.73]

warmup:   7%|▋         | 49/750 [00:23<04:53,  2.39it/s, 511 steps of size 2.12e-03. acc. prob=0.73]

warmup:   7%|▋         | 50/750 [00:23<05:01,  2.32it/s, 255 steps of size 3.47e-03. acc. prob=0.73]

warmup:   7%|▋         | 51/750 [00:24<04:27,  2.61it/s, 127 steps of size 1.92e-03. acc. prob=0.73]

warmup:   7%|▋         | 52/750 [00:24<04:40,  2.49it/s, 255 steps of size 3.06e-03. acc. prob=0.73]

warmup:   7%|▋         | 53/750 [00:24<04:08,  2.80it/s, 130 steps of size 4.51e-03. acc. prob=0.74]

warmup:   7%|▋         | 54/750 [00:25<03:26,  3.37it/s, 77 steps of size 1.79e-03. acc. prob=0.73]

warmup:   7%|▋         | 55/750 [00:25<03:57,  2.93it/s, 255 steps of size 2.91e-03. acc. prob=0.74]

warmup:   7%|▋         | 56/750 [00:25<03:34,  3.24it/s, 127 steps of size 3.14e-03. acc. prob=0.74]

warmup:   8%|▊         | 57/750 [00:25<03:17,  3.50it/s, 127 steps of size 5.00e-03. acc. prob=0.74]

warmup:   8%|▊         | 59/750 [00:26<04:13,  2.72it/s, 511 steps of size 1.76e-03. acc. prob=0.73]

warmup:   8%|▊         | 60/750 [00:27<04:28,  2.57it/s, 255 steps of size 2.88e-03. acc. prob=0.74]

warmup:   8%|▊         | 61/750 [00:27<03:39,  3.15it/s, 63 steps of size 2.12e-03. acc. prob=0.74]

warmup:   8%|▊         | 62/750 [00:27<04:04,  2.81it/s, 255 steps of size 2.85e-03. acc. prob=0.74]

warmup:   8%|▊         | 63/750 [00:28<03:39,  3.12it/s, 127 steps of size 4.61e-03. acc. prob=0.74]

warmup:   9%|▊         | 64/750 [00:28<03:21,  3.40it/s, 127 steps of size 4.65e-03. acc. prob=0.74]

warmup:   9%|▉         | 66/750 [00:29<04:21,  2.62it/s, 511 steps of size 1.82e-03. acc. prob=0.74]

warmup:   9%|▉         | 67/750 [00:29<04:33,  2.50it/s, 255 steps of size 2.23e-03. acc. prob=0.74]

warmup:   9%|▉         | 68/750 [00:30<04:42,  2.41it/s, 255 steps of size 3.31e-03. acc. prob=0.74]

warmup:   9%|▉         | 69/750 [00:30<04:08,  2.74it/s, 127 steps of size 5.00e-03. acc. prob=0.75]

warmup:   9%|▉         | 71/750 [00:31<04:36,  2.46it/s, 511 steps of size 1.59e-03. acc. prob=0.74]

warmup:  10%|▉         | 72/750 [00:31<04:44,  2.39it/s, 255 steps of size 2.51e-03. acc. prob=0.74]

warmup:  10%|▉         | 73/750 [00:32<04:50,  2.33it/s, 255 steps of size 3.81e-03. acc. prob=0.75]

warmup:  10%|▉         | 74/750 [00:32<04:13,  2.66it/s, 127 steps of size 1.51e-03. acc. prob=0.74]

warmup:  10%|█         | 75/750 [00:33<04:31,  2.49it/s, 255 steps of size 1.39e-03. acc. prob=0.74]

warmup:  10%|█         | 76/750 [00:33<06:08,  1.83it/s, 511 steps of size 2.17e-03. acc. prob=0.75]

warmup:  10%|█         | 77/750 [00:34<05:50,  1.92it/s, 255 steps of size 3.37e-03. acc. prob=0.75]

warmup:  10%|█         | 78/750 [00:34<04:52,  2.30it/s, 127 steps of size 3.21e-03. acc. prob=0.75]

warmup:  11%|█         | 79/750 [00:34<04:11,  2.67it/s, 127 steps of size 3.80e-03. acc. prob=0.75]

warmup:  11%|█         | 80/750 [00:35<03:41,  3.02it/s, 127 steps of size 8.06e-04. acc. prob=0.74]

warmup:  11%|█         | 81/750 [00:36<05:37,  1.98it/s, 511 steps of size 1.18e-03. acc. prob=0.74]

warmup:  11%|█         | 82/750 [00:36<07:10,  1.55it/s, 511 steps of size 1.52e-03. acc. prob=0.75]

warmup:  11%|█         | 83/750 [00:37<06:32,  1.70it/s, 255 steps of size 2.34e-03. acc. prob=0.75]

warmup:  11%|█         | 84/750 [00:37<05:35,  1.98it/s, 169 steps of size 3.23e-03. acc. prob=0.75]

warmup:  11%|█▏        | 86/750 [00:38<04:13,  2.62it/s, 255 steps of size 1.56e-03. acc. prob=0.75]

warmup:  12%|█▏        | 87/750 [00:38<04:27,  2.47it/s, 255 steps of size 2.06e-03. acc. prob=0.75]

warmup:  12%|█▏        | 88/750 [00:39<04:36,  2.39it/s, 255 steps of size 8.84e-04. acc. prob=0.74]

warmup:  12%|█▏        | 89/750 [00:40<06:08,  1.80it/s, 511 steps of size 1.33e-03. acc. prob=0.75]

warmup:  12%|█▏        | 90/750 [00:40<05:48,  1.89it/s, 255 steps of size 9.34e-04. acc. prob=0.75]

warmup:  12%|█▏        | 91/750 [00:41<07:00,  1.57it/s, 511 steps of size 1.17e-03. acc. prob=0.75]

warmup:  12%|█▏        | 92/750 [00:42<09:19,  1.18it/s, 767 steps of size 1.25e-03. acc. prob=0.75]

warmup:  12%|█▏        | 93/750 [00:43<09:30,  1.15it/s, 511 steps of size 1.42e-03. acc. prob=0.75]

warmup:  13%|█▎        | 94/750 [00:44<08:10,  1.34it/s, 255 steps of size 1.59e-03. acc. prob=0.75]

warmup:  13%|█▎        | 95/750 [00:44<07:16,  1.50it/s, 255 steps of size 2.17e-03. acc. prob=0.75]

warmup:  13%|█▎        | 96/750 [00:45<06:29,  1.68it/s, 239 steps of size 2.53e-03. acc. prob=0.75]

warmup:  13%|█▎        | 97/750 [00:45<05:17,  2.06it/s, 127 steps of size 6.42e-04. acc. prob=0.75]

warmup:  13%|█▎        | 98/750 [00:46<06:40,  1.63it/s, 511 steps of size 9.53e-04. acc. prob=0.75]

warmup:  13%|█▎        | 99/750 [00:47<07:40,  1.41it/s, 511 steps of size 1.43e-03. acc. prob=0.75]

warmup:  13%|█▎        | 100/750 [00:48<08:20,  1.30it/s, 511 steps of size 1.91e-03. acc. prob=0.75]

warmup:  13%|█▎        | 101/750 [00:49<11:49,  1.09s/it, 1023 steps of size 2.75e-02. acc. prob=0.76]

warmup:  14%|█▎        | 102/750 [00:50<09:00,  1.20it/s, 127 steps of size 4.85e-02. acc. prob=0.76]

warmup:  14%|█▎        | 103/750 [00:50<07:02,  1.53it/s, 127 steps of size 8.81e-02. acc. prob=0.76]

warmup:  14%|█▍        | 104/750 [00:50<05:17,  2.04it/s, 63 steps of size 1.66e-01. acc. prob=0.76]

warmup:  14%|█▍        | 106/750 [00:50<03:06,  3.44it/s, 31 steps of size 4.85e-02. acc. prob=0.76]

warmup:  14%|█▍        | 107/750 [00:50<02:56,  3.64it/s, 127 steps of size 8.99e-02. acc. prob=0.76]

warmup:  14%|█▍        | 108/750 [00:50<02:29,  4.29it/s, 63 steps of size 1.64e-01. acc. prob=0.76]

warmup:  15%|█▍        | 110/750 [00:51<01:40,  6.35it/s, 31 steps of size 1.46e-01. acc. prob=0.76]

warmup:  15%|█▍        | 112/750 [00:51<01:36,  6.60it/s, 127 steps of size 6.71e-02. acc. prob=0.76]

warmup:  15%|█▌        | 113/750 [00:51<01:31,  6.97it/s, 63 steps of size 1.15e-01. acc. prob=0.76]

warmup:  15%|█▌        | 115/750 [00:51<01:10,  9.07it/s, 31 steps of size 1.93e-01. acc. prob=0.77]

warmup:  16%|█▌        | 117/750 [00:51<01:04,  9.88it/s, 63 steps of size 1.67e-01. acc. prob=0.77]

warmup:  16%|█▌        | 119/750 [00:51<00:53, 11.78it/s, 31 steps of size 1.83e-01. acc. prob=0.77]

warmup:  16%|█▌        | 121/750 [00:52<00:55, 11.41it/s, 63 steps of size 1.06e-01. acc. prob=0.77]

warmup:  16%|█▋        | 123/750 [00:52<00:48, 12.85it/s, 31 steps of size 2.62e-01. acc. prob=0.77]

warmup:  17%|█▋        | 125/750 [00:52<00:56, 11.09it/s, 127 steps of size 8.31e-02. acc. prob=0.77]

warmup:  17%|█▋        | 126/750 [00:52<00:58, 10.60it/s, 63 steps of size 3.60e-02. acc. prob=0.76]

warmup:  17%|█▋        | 127/750 [00:52<01:15,  8.21it/s, 127 steps of size 6.25e-02. acc. prob=0.77]

warmup:  17%|█▋        | 128/750 [00:52<01:14,  8.34it/s, 63 steps of size 1.07e-01. acc. prob=0.77]

warmup:  17%|█▋        | 129/750 [00:52<01:13,  8.44it/s, 63 steps of size 1.68e-01. acc. prob=0.77]

warmup:  17%|█▋        | 131/750 [00:53<01:19,  7.80it/s, 127 steps of size 6.47e-02. acc. prob=0.77]

warmup:  18%|█▊        | 132/750 [00:53<01:17,  8.01it/s, 63 steps of size 1.10e-01. acc. prob=0.77]

warmup:  18%|█▊        | 134/750 [00:53<01:00, 10.24it/s, 31 steps of size 1.72e-01. acc. prob=0.77]

warmup:  18%|█▊        | 138/750 [00:53<00:45, 13.58it/s, 63 steps of size 9.01e-02. acc. prob=0.77]

warmup:  19%|█▊        | 140/750 [00:53<00:41, 14.54it/s, 31 steps of size 1.23e-01. acc. prob=0.77]

warmup:  19%|█▉        | 141/750 [00:53<00:46, 13.14it/s, 63 steps of size 1.41e-01. acc. prob=0.77]

warmup:  19%|█▉        | 143/750 [00:54<00:42, 14.36it/s, 31 steps of size 2.25e-01. acc. prob=0.77]

warmup:  19%|█▉        | 145/750 [00:54<00:44, 13.46it/s, 63 steps of size 7.79e-02. acc. prob=0.77]

warmup:  19%|█▉        | 146/750 [00:54<00:49, 12.28it/s, 63 steps of size 1.17e-01. acc. prob=0.77]

warmup:  20%|█▉        | 148/750 [00:54<00:50, 11.98it/s, 63 steps of size 4.58e-02. acc. prob=0.77]

warmup:  20%|█▉        | 149/750 [00:54<01:07,  8.92it/s, 127 steps of size 6.87e-02. acc. prob=0.77]

warmup:  20%|██        | 150/750 [00:54<01:07,  8.90it/s, 63 steps of size 9.29e-02. acc. prob=0.77]

warmup:  20%|██        | 151/750 [00:54<01:07,  8.88it/s, 63 steps of size 1.33e+00. acc. prob=0.77]

warmup:  21%|██        | 154/750 [00:55<00:57, 10.42it/s, 127 steps of size 3.42e-02. acc. prob=0.77]

warmup:  21%|██        | 155/750 [00:55<01:12,  8.23it/s, 127 steps of size 5.32e-02. acc. prob=0.77]

warmup:  21%|██        | 156/750 [00:55<01:11,  8.35it/s, 63 steps of size 7.84e-02. acc. prob=0.77]

warmup:  21%|██        | 157/750 [00:55<01:10,  8.44it/s, 63 steps of size 1.34e-01. acc. prob=0.77]

warmup:  21%|██        | 159/750 [00:56<01:37,  6.06it/s, 255 steps of size 3.30e-02. acc. prob=0.77]

warmup:  21%|██▏       | 160/750 [00:56<01:45,  5.58it/s, 127 steps of size 5.00e-02. acc. prob=0.77]

warmup:  21%|██▏       | 161/750 [00:56<01:52,  5.24it/s, 127 steps of size 8.87e-02. acc. prob=0.77]

warmup:  22%|██▏       | 163/750 [00:56<01:19,  7.37it/s, 31 steps of size 1.98e-01. acc. prob=0.77]

warmup:  22%|██▏       | 167/750 [00:56<00:49, 11.78it/s, 47 steps of size 2.29e-01. acc. prob=0.77]

warmup:  23%|██▎       | 170/750 [00:57<00:45, 12.71it/s, 95 steps of size 1.35e-01. acc. prob=0.77]

warmup:  23%|██▎       | 173/750 [00:57<00:51, 11.17it/s, 127 steps of size 5.47e-02. acc. prob=0.77]

warmup:  23%|██▎       | 174/750 [00:57<00:53, 10.75it/s, 63 steps of size 9.20e-02. acc. prob=0.77]

warmup:  23%|██▎       | 175/750 [00:57<00:55, 10.34it/s, 63 steps of size 6.98e-02. acc. prob=0.77]

warmup:  23%|██▎       | 176/750 [00:57<00:57, 10.01it/s, 63 steps of size 1.20e-01. acc. prob=0.77]

warmup:  24%|██▍       | 179/750 [00:57<00:49, 11.53it/s, 63 steps of size 1.13e-01. acc. prob=0.77]

warmup:  24%|██▍       | 181/750 [00:58<00:44, 12.92it/s, 31 steps of size 1.63e-01. acc. prob=0.77]

warmup:  24%|██▍       | 183/750 [00:58<00:54, 10.33it/s, 127 steps of size 3.81e-02. acc. prob=0.77]

warmup:  25%|██▍       | 184/750 [00:58<01:07,  8.34it/s, 127 steps of size 6.21e-02. acc. prob=0.77]

warmup:  25%|██▍       | 185/750 [00:58<01:20,  7.04it/s, 127 steps of size 1.07e-01. acc. prob=0.77]

warmup:  25%|██▍       | 187/750 [00:58<01:02,  9.04it/s, 31 steps of size 2.56e-01. acc. prob=0.78]

warmup:  25%|██▌       | 189/750 [00:59<01:05,  8.57it/s, 127 steps of size 7.33e-02. acc. prob=0.77]

warmup:  25%|██▌       | 190/750 [00:59<01:05,  8.61it/s, 63 steps of size 8.12e-02. acc. prob=0.77]

warmup:  25%|██▌       | 191/750 [00:59<01:04,  8.65it/s, 63 steps of size 9.98e-02. acc. prob=0.77]

warmup:  26%|██▌       | 194/750 [00:59<00:59,  9.39it/s, 127 steps of size 5.37e-02. acc. prob=0.77]

warmup:  26%|██▌       | 195/750 [00:59<00:59,  9.27it/s, 63 steps of size 8.48e-02. acc. prob=0.77]

warmup:  26%|██▌       | 196/750 [00:59<01:00,  9.17it/s, 63 steps of size 1.56e-02. acc. prob=0.77]

warmup:  26%|██▋       | 197/750 [01:00<01:45,  5.26it/s, 255 steps of size 2.60e-02. acc. prob=0.77]

warmup:  26%|██▋       | 198/750 [01:00<01:52,  4.91it/s, 127 steps of size 4.27e-02. acc. prob=0.77]

warmup:  27%|██▋       | 199/750 [01:00<01:55,  4.75it/s, 127 steps of size 6.01e-02. acc. prob=0.77]

warmup:  27%|██▋       | 200/750 [01:00<01:41,  5.43it/s, 63 steps of size 7.37e-02. acc. prob=0.77]

warmup:  27%|██▋       | 201/750 [01:01<01:30,  6.08it/s, 63 steps of size 1.14e-01. acc. prob=0.78]

warmup:  27%|██▋       | 203/750 [01:01<01:10,  7.76it/s, 63 steps of size 1.03e-01. acc. prob=0.78]

warmup:  27%|██▋       | 204/750 [01:01<01:08,  7.99it/s, 63 steps of size 1.47e-01. acc. prob=0.78]

warmup:  27%|██▋       | 206/750 [01:01<01:06,  8.18it/s, 127 steps of size 5.50e-02. acc. prob=0.77]

warmup:  28%|██▊       | 207/750 [01:01<01:05,  8.32it/s, 63 steps of size 8.73e-02. acc. prob=0.78]

warmup:  28%|██▊       | 210/750 [01:01<00:46, 11.60it/s, 31 steps of size 7.14e-02. acc. prob=0.78]

warmup:  28%|██▊       | 211/750 [01:01<00:49, 10.95it/s, 63 steps of size 1.12e-01. acc. prob=0.78]

warmup:  28%|██▊       | 213/750 [01:02<00:58,  9.18it/s, 127 steps of size 3.80e-02. acc. prob=0.77]

warmup:  29%|██▊       | 214/750 [01:02<01:11,  7.53it/s, 127 steps of size 5.96e-02. acc. prob=0.78]

warmup:  29%|██▊       | 215/750 [01:02<01:08,  7.77it/s, 63 steps of size 8.77e-02. acc. prob=0.78]

warmup:  29%|██▉       | 217/750 [01:02<00:53,  9.88it/s, 31 steps of size 2.51e-02. acc. prob=0.77]

warmup:  29%|██▉       | 218/750 [01:02<01:09,  7.71it/s, 127 steps of size 3.89e-02. acc. prob=0.77]

warmup:  29%|██▉       | 219/750 [01:03<01:21,  6.52it/s, 127 steps of size 6.03e-02. acc. prob=0.78]

warmup:  29%|██▉       | 220/750 [01:03<01:15,  6.98it/s, 63 steps of size 8.36e-02. acc. prob=0.78]

warmup:  29%|██▉       | 221/750 [01:03<01:11,  7.36it/s, 63 steps of size 1.24e-01. acc. prob=0.78]

warmup:  30%|██▉       | 223/750 [01:03<00:53,  9.86it/s, 31 steps of size 2.15e-01. acc. prob=0.78]

warmup:  30%|███       | 226/750 [01:03<00:39, 13.27it/s, 31 steps of size 9.01e-02. acc. prob=0.78]

warmup:  30%|███       | 228/750 [01:03<00:36, 14.36it/s, 31 steps of size 1.54e-01. acc. prob=0.78]

warmup:  31%|███       | 230/750 [01:03<00:39, 13.29it/s, 63 steps of size 6.52e-02. acc. prob=0.78]

warmup:  31%|███       | 231/750 [01:04<00:45, 11.50it/s, 63 steps of size 8.86e-02. acc. prob=0.78]

warmup:  31%|███       | 233/750 [01:04<00:39, 13.03it/s, 31 steps of size 1.25e-01. acc. prob=0.78]

warmup:  31%|███▏      | 235/750 [01:04<00:48, 10.63it/s, 127 steps of size 3.81e-02. acc. prob=0.78]

warmup:  31%|███▏      | 236/750 [01:04<01:01,  8.38it/s, 127 steps of size 5.49e-02. acc. prob=0.78]

warmup:  32%|███▏      | 237/750 [01:04<01:00,  8.47it/s, 63 steps of size 5.18e-02. acc. prob=0.78]

warmup:  32%|███▏      | 238/750 [01:04<00:59,  8.55it/s, 63 steps of size 2.73e-02. acc. prob=0.78]

warmup:  32%|███▏      | 239/750 [01:05<01:13,  6.91it/s, 127 steps of size 3.65e-02. acc. prob=0.78]

warmup:  32%|███▏      | 240/750 [01:05<01:09,  7.32it/s, 63 steps of size 5.21e-02. acc. prob=0.78]

warmup:  32%|███▏      | 241/750 [01:05<01:06,  7.65it/s, 63 steps of size 5.79e-02. acc. prob=0.78]

warmup:  32%|███▏      | 242/750 [01:05<01:04,  7.89it/s, 63 steps of size 4.20e-02. acc. prob=0.78]

warmup:  32%|███▏      | 243/750 [01:05<01:19,  6.37it/s, 127 steps of size 6.01e-02. acc. prob=0.78]

warmup:  33%|███▎      | 244/750 [01:05<01:13,  6.90it/s, 63 steps of size 8.51e-02. acc. prob=0.78]

warmup:  33%|███▎      | 245/750 [01:05<01:08,  7.38it/s, 63 steps of size 1.24e-01. acc. prob=0.78]

warmup:  33%|███▎      | 246/750 [01:06<01:05,  7.73it/s, 63 steps of size 4.46e-02. acc. prob=0.78]

warmup:  33%|███▎      | 247/750 [01:06<01:02,  8.03it/s, 63 steps of size 6.13e-02. acc. prob=0.78]

warmup:  33%|███▎      | 248/750 [01:06<01:00,  8.25it/s, 63 steps of size 6.98e-02. acc. prob=0.78]

warmup:  33%|███▎      | 249/750 [01:06<00:59,  8.37it/s, 63 steps of size 1.00e-01. acc. prob=0.78]

warmup:  33%|███▎      | 251/750 [01:06<00:44, 11.09it/s, 31 steps of size 1.92e+00. acc. prob=0.78]

warmup:  34%|███▍      | 254/750 [01:06<00:32, 15.14it/s, 63 steps of size 4.50e-02. acc. prob=0.78]

warmup:  34%|███▍      | 255/750 [01:06<00:37, 13.33it/s, 63 steps of size 6.69e-02. acc. prob=0.78]

warmup:  34%|███▍      | 256/750 [01:06<00:41, 12.01it/s, 63 steps of size 1.13e-01. acc. prob=0.78]

warmup:  34%|███▍      | 258/750 [01:06<00:39, 12.34it/s, 63 steps of size 6.01e-02. acc. prob=0.78]

warmup:  35%|███▍      | 259/750 [01:07<00:45, 10.90it/s, 63 steps of size 9.53e-02. acc. prob=0.78]

warmup:  35%|███▍      | 260/750 [01:07<00:47, 10.33it/s, 63 steps of size 3.15e-02. acc. prob=0.78]

warmup:  35%|███▍      | 261/750 [01:07<01:03,  7.73it/s, 127 steps of size 5.57e-02. acc. prob=0.78]

warmup:  35%|███▍      | 262/750 [01:07<01:30,  5.40it/s, 191 steps of size 1.04e-01. acc. prob=0.78]

warmup:  35%|███▌      | 264/750 [01:07<01:08,  7.04it/s, 63 steps of size 8.43e-02. acc. prob=0.78]

warmup:  35%|███▌      | 265/750 [01:08<01:05,  7.39it/s, 63 steps of size 1.57e-01. acc. prob=0.78]

warmup:  36%|███▌      | 267/750 [01:08<00:55,  8.70it/s, 63 steps of size 1.29e-01. acc. prob=0.78]

warmup:  36%|███▌      | 270/750 [01:08<00:44, 10.90it/s, 63 steps of size 1.02e-01. acc. prob=0.78]

warmup:  36%|███▌      | 271/750 [01:08<00:50,  9.45it/s, 95 steps of size 6.43e-02. acc. prob=0.78]

warmup:  36%|███▋      | 272/750 [01:08<00:51,  9.32it/s, 63 steps of size 1.18e-01. acc. prob=0.78]

warmup:  37%|███▋      | 274/750 [01:09<01:17,  6.17it/s, 255 steps of size 2.93e-02. acc. prob=0.78]

warmup:  37%|███▋      | 275/750 [01:09<01:23,  5.68it/s, 127 steps of size 5.33e-02. acc. prob=0.78]

warmup:  37%|███▋      | 276/750 [01:09<01:16,  6.18it/s, 63 steps of size 9.35e-02. acc. prob=0.78]

warmup:  37%|███▋      | 277/750 [01:09<01:10,  6.68it/s, 63 steps of size 1.45e-01. acc. prob=0.78]

warmup:  37%|███▋      | 279/750 [01:09<00:52,  9.01it/s, 31 steps of size 1.10e-01. acc. prob=0.78]

warmup:  38%|███▊      | 282/750 [01:09<00:35, 13.34it/s, 11 steps of size 4.80e-02. acc. prob=0.78]

warmup:  38%|███▊      | 283/750 [01:10<00:38, 12.20it/s, 63 steps of size 8.18e-02. acc. prob=0.78]

warmup:  38%|███▊      | 284/750 [01:10<00:41, 11.30it/s, 63 steps of size 8.95e-02. acc. prob=0.78]

warmup:  38%|███▊      | 287/750 [01:10<00:43, 10.76it/s, 127 steps of size 4.83e-02. acc. prob=0.78]

warmup:  38%|███▊      | 288/750 [01:10<00:44, 10.35it/s, 63 steps of size 6.08e-02. acc. prob=0.78]

warmup:  39%|███▊      | 289/750 [01:10<00:46, 10.00it/s, 63 steps of size 1.85e-02. acc. prob=0.78]

warmup:  39%|███▊      | 290/750 [01:11<01:23,  5.50it/s, 255 steps of size 3.16e-02. acc. prob=0.78]

warmup:  39%|███▉      | 291/750 [01:11<01:28,  5.18it/s, 127 steps of size 4.38e-02. acc. prob=0.78]

warmup:  39%|███▉      | 292/750 [01:11<01:32,  4.93it/s, 127 steps of size 4.91e-02. acc. prob=0.78]

warmup:  39%|███▉      | 293/750 [01:11<01:21,  5.60it/s, 63 steps of size 2.18e-02. acc. prob=0.78]

warmup:  39%|███▉      | 294/750 [01:11<01:27,  5.20it/s, 127 steps of size 3.46e-02. acc. prob=0.78]

warmup:  39%|███▉      | 295/750 [01:12<01:32,  4.94it/s, 127 steps of size 5.66e-02. acc. prob=0.78]

warmup:  39%|███▉      | 296/750 [01:12<01:20,  5.66it/s, 63 steps of size 6.41e-02. acc. prob=0.78]

warmup:  40%|███▉      | 297/750 [01:12<01:11,  6.32it/s, 63 steps of size 5.90e-02. acc. prob=0.78]

warmup:  40%|███▉      | 298/750 [01:12<01:07,  6.67it/s, 63 steps of size 3.08e-02. acc. prob=0.78]

warmup:  40%|███▉      | 299/750 [01:12<01:18,  5.76it/s, 127 steps of size 4.89e-02. acc. prob=0.78]

warmup:  40%|████      | 301/750 [01:13<01:10,  6.37it/s, 127 steps of size 5.03e-02. acc. prob=0.78]

warmup:  40%|████      | 303/750 [01:13<02:03,  3.62it/s, 511 steps of size 1.27e-02. acc. prob=0.78]

warmup:  41%|████      | 304/750 [01:14<02:21,  3.16it/s, 255 steps of size 2.02e-02. acc. prob=0.78]

warmup:  41%|████      | 305/750 [01:14<02:35,  2.86it/s, 255 steps of size 3.17e-02. acc. prob=0.78]

warmup:  41%|████      | 306/750 [01:15<02:21,  3.13it/s, 127 steps of size 4.95e-02. acc. prob=0.78]

warmup:  41%|████      | 307/750 [01:15<01:57,  3.78it/s, 63 steps of size 5.69e-02. acc. prob=0.78]

warmup:  41%|████      | 308/750 [01:15<01:38,  4.49it/s, 63 steps of size 2.15e-02. acc. prob=0.78]

warmup:  41%|████      | 309/750 [01:15<02:07,  3.46it/s, 255 steps of size 3.34e-02. acc. prob=0.78]

warmup:  41%|████▏     | 310/750 [01:15<01:59,  3.68it/s, 127 steps of size 4.37e-02. acc. prob=0.78]

warmup:  41%|████▏     | 311/750 [01:16<01:53,  3.86it/s, 127 steps of size 6.60e-02. acc. prob=0.78]

warmup:  42%|████▏     | 312/750 [01:16<01:34,  4.62it/s, 63 steps of size 1.20e-02. acc. prob=0.78]

warmup:  42%|████▏     | 313/750 [01:16<02:07,  3.42it/s, 255 steps of size 1.88e-02. acc. prob=0.78]

warmup:  42%|████▏     | 314/750 [01:17<02:28,  2.93it/s, 255 steps of size 2.98e-02. acc. prob=0.78]

warmup:  42%|████▏     | 315/750 [01:17<02:13,  3.26it/s, 127 steps of size 4.63e-02. acc. prob=0.78]

warmup:  42%|████▏     | 316/750 [01:17<01:48,  4.01it/s, 63 steps of size 6.63e-02. acc. prob=0.78]

warmup:  42%|████▏     | 317/750 [01:17<01:30,  4.79it/s, 63 steps of size 1.02e-01. acc. prob=0.78]

warmup:  43%|████▎     | 319/750 [01:17<00:59,  7.22it/s, 31 steps of size 8.40e-02. acc. prob=0.78]

warmup:  43%|████▎     | 320/750 [01:17<00:56,  7.55it/s, 63 steps of size 1.29e-01. acc. prob=0.78]

warmup:  43%|████▎     | 322/750 [01:18<00:43,  9.89it/s, 31 steps of size 9.78e-02. acc. prob=0.78]

warmup:  43%|████▎     | 324/750 [01:18<00:35, 11.85it/s, 31 steps of size 1.07e-01. acc. prob=0.78]

warmup:  43%|████▎     | 326/750 [01:18<00:44,  9.58it/s, 127 steps of size 5.36e-02. acc. prob=0.78]

warmup:  44%|████▎     | 327/750 [01:18<00:44,  9.42it/s, 63 steps of size 8.01e-02. acc. prob=0.78]

warmup:  44%|████▎     | 328/750 [01:18<00:45,  9.29it/s, 63 steps of size 9.51e-02. acc. prob=0.78]

warmup:  44%|████▍     | 329/750 [01:18<00:45,  9.17it/s, 63 steps of size 1.32e-01. acc. prob=0.78]

warmup:  44%|████▍     | 331/750 [01:18<00:36, 11.44it/s, 31 steps of size 1.76e-01. acc. prob=0.78]

warmup:  45%|████▍     | 334/750 [01:19<00:26, 15.64it/s, 15 steps of size 2.25e-01. acc. prob=0.78]

warmup:  45%|████▍     | 337/750 [01:19<00:25, 16.49it/s, 63 steps of size 1.19e-01. acc. prob=0.78]

warmup:  45%|████▌     | 339/750 [01:19<00:24, 16.86it/s, 31 steps of size 8.66e-02. acc. prob=0.78]

warmup:  45%|████▌     | 341/750 [01:19<00:23, 17.14it/s, 31 steps of size 6.54e-02. acc. prob=0.78]

warmup:  46%|████▌     | 342/750 [01:19<00:27, 14.81it/s, 63 steps of size 9.54e-02. acc. prob=0.78]

warmup:  46%|████▌     | 343/750 [01:19<00:31, 13.09it/s, 63 steps of size 1.17e-01. acc. prob=0.78]

warmup:  46%|████▌     | 345/750 [01:19<00:32, 12.61it/s, 63 steps of size 1.25e-01. acc. prob=0.78]

warmup:  46%|████▋     | 347/750 [01:19<00:28, 13.98it/s, 31 steps of size 1.66e-01. acc. prob=0.78]

warmup:  47%|████▋     | 351/750 [01:20<00:23, 17.19it/s, 63 steps of size 6.19e-02. acc. prob=0.78]

warmup:  47%|████▋     | 352/750 [01:20<00:26, 15.01it/s, 63 steps of size 8.83e-02. acc. prob=0.78]

warmup:  47%|████▋     | 354/750 [01:20<00:28, 13.81it/s, 63 steps of size 1.54e-01. acc. prob=0.78]

warmup:  47%|████▋     | 356/750 [01:20<00:26, 14.76it/s, 31 steps of size 9.14e-02. acc. prob=0.78]

warmup:  48%|████▊     | 357/750 [01:20<00:33, 11.63it/s, 95 steps of size 7.41e-02. acc. prob=0.78]

warmup:  48%|████▊     | 358/750 [01:20<00:35, 10.93it/s, 63 steps of size 1.06e-01. acc. prob=0.78]

warmup:  48%|████▊     | 360/750 [01:20<00:30, 12.74it/s, 31 steps of size 1.60e-01. acc. prob=0.78]

warmup:  48%|████▊     | 363/750 [01:20<00:23, 16.60it/s, 31 steps of size 1.40e-01. acc. prob=0.78]

warmup:  49%|████▊     | 365/750 [01:21<00:22, 16.96it/s, 31 steps of size 7.33e-02. acc. prob=0.78]

warmup:  49%|████▉     | 366/750 [01:21<00:26, 14.63it/s, 63 steps of size 1.01e-01. acc. prob=0.78]

warmup:  49%|████▉     | 369/750 [01:21<00:26, 14.44it/s, 63 steps of size 4.33e-02. acc. prob=0.78]

warmup:  49%|████▉     | 370/750 [01:21<00:36, 10.42it/s, 127 steps of size 6.05e-02. acc. prob=0.78]

warmup:  49%|████▉     | 371/750 [01:21<00:37, 10.07it/s, 63 steps of size 8.43e-02. acc. prob=0.78]

warmup:  50%|████▉     | 373/750 [01:21<00:35, 10.68it/s, 63 steps of size 3.09e-02. acc. prob=0.78]

warmup:  50%|████▉     | 374/750 [01:22<00:45,  8.28it/s, 127 steps of size 3.72e-02. acc. prob=0.78]

warmup:  50%|█████     | 375/750 [01:22<00:54,  6.90it/s, 127 steps of size 4.55e-02. acc. prob=0.78]

warmup:  50%|█████     | 376/750 [01:22<01:01,  6.06it/s, 127 steps of size 6.25e-02. acc. prob=0.78]

warmup:  50%|█████     | 377/750 [01:22<00:56,  6.59it/s, 63 steps of size 8.44e-02. acc. prob=0.78]

warmup:  50%|█████     | 378/750 [01:22<00:52,  7.07it/s, 63 steps of size 1.17e-01. acc. prob=0.78]

warmup:  51%|█████     | 380/750 [01:22<00:38,  9.60it/s, 31 steps of size 1.67e-01. acc. prob=0.78]

warmup:  51%|█████     | 382/750 [01:23<00:33, 10.96it/s, 63 steps of size 1.23e-01. acc. prob=0.78]

warmup:  51%|█████     | 384/750 [01:23<00:28, 12.71it/s, 31 steps of size 7.08e-02. acc. prob=0.78]

warmup:  51%|█████▏    | 385/750 [01:23<00:31, 11.64it/s, 63 steps of size 9.38e-02. acc. prob=0.78]

warmup:  52%|█████▏    | 387/750 [01:23<00:27, 13.25it/s, 31 steps of size 1.21e-01. acc. prob=0.78]

warmup:  52%|█████▏    | 389/750 [01:23<00:28, 12.72it/s, 63 steps of size 1.27e-01. acc. prob=0.78]

warmup:  52%|█████▏    | 391/750 [01:23<00:25, 13.95it/s, 31 steps of size 1.46e-01. acc. prob=0.78]

warmup:  52%|█████▏    | 393/750 [01:23<00:24, 14.77it/s, 31 steps of size 1.70e-01. acc. prob=0.78]

warmup:  53%|█████▎    | 395/750 [01:23<00:23, 15.07it/s, 31 steps of size 1.45e-01. acc. prob=0.78]

warmup:  53%|█████▎    | 397/750 [01:24<00:22, 15.67it/s, 31 steps of size 2.16e-01. acc. prob=0.78]

warmup:  53%|█████▎    | 399/750 [01:24<00:24, 14.62it/s, 63 steps of size 8.58e-02. acc. prob=0.78]

warmup:  53%|█████▎    | 400/750 [01:24<00:26, 13.03it/s, 63 steps of size 1.01e-01. acc. prob=0.78]

warmup:  54%|█████▎    | 402/750 [01:24<00:24, 14.30it/s, 31 steps of size 6.73e-02. acc. prob=0.78]

warmup:  54%|█████▎    | 403/750 [01:24<00:27, 12.75it/s, 63 steps of size 9.03e-02. acc. prob=0.78]

warmup:  54%|█████▍    | 404/750 [01:24<00:29, 11.62it/s, 63 steps of size 1.15e-01. acc. prob=0.78]

warmup:  54%|█████▍    | 406/750 [01:24<00:25, 13.38it/s, 31 steps of size 7.04e-02. acc. prob=0.78]

warmup:  54%|█████▍    | 407/750 [01:24<00:28, 12.04it/s, 63 steps of size 8.72e-02. acc. prob=0.78]

warmup:  54%|█████▍    | 408/750 [01:25<00:30, 11.08it/s, 63 steps of size 7.13e-02. acc. prob=0.78]

warmup:  55%|█████▍    | 409/750 [01:25<00:32, 10.42it/s, 63 steps of size 9.41e-02. acc. prob=0.78]

warmup:  55%|█████▍    | 410/750 [01:25<00:34,  9.93it/s, 63 steps of size 1.08e-01. acc. prob=0.78]

warmup:  55%|█████▌    | 413/750 [01:25<00:24, 13.87it/s, 31 steps of size 1.44e-01. acc. prob=0.78]

warmup:  55%|█████▌    | 415/750 [01:25<00:25, 13.09it/s, 63 steps of size 6.85e-02. acc. prob=0.78]

warmup:  55%|█████▌    | 416/750 [01:25<00:27, 11.93it/s, 63 steps of size 9.14e-02. acc. prob=0.78]

warmup:  56%|█████▌    | 417/750 [01:25<00:30, 11.04it/s, 63 steps of size 1.23e-01. acc. prob=0.78]

warmup:  56%|█████▌    | 419/750 [01:25<00:25, 12.92it/s, 31 steps of size 1.77e-01. acc. prob=0.78]

warmup:  56%|█████▌    | 421/750 [01:26<00:23, 14.20it/s, 31 steps of size 1.88e-01. acc. prob=0.78]

warmup:  56%|█████▋    | 423/750 [01:26<00:21, 15.22it/s, 31 steps of size 1.41e-01. acc. prob=0.78]

warmup:  57%|█████▋    | 425/750 [01:26<00:20, 15.92it/s, 31 steps of size 1.31e-01. acc. prob=0.78]

warmup:  57%|█████▋    | 427/750 [01:26<00:19, 16.48it/s, 31 steps of size 2.12e-01. acc. prob=0.79]

warmup:  57%|█████▋    | 431/750 [01:26<00:14, 22.34it/s, 15 steps of size 3.13e-01. acc. prob=0.79]

warmup:  58%|█████▊    | 435/750 [01:26<00:11, 26.55it/s, 15 steps of size 1.20e-01. acc. prob=0.78]

warmup:  58%|█████▊    | 437/750 [01:26<00:13, 23.94it/s, 31 steps of size 2.07e-01. acc. prob=0.79]

warmup:  59%|█████▉    | 441/750 [01:26<00:11, 27.71it/s, 15 steps of size 3.22e-01. acc. prob=0.79]

warmup:  59%|█████▉    | 444/750 [01:26<00:11, 26.44it/s, 31 steps of size 1.69e-01. acc. prob=0.79]

warmup:  59%|█████▉    | 446/750 [01:27<00:12, 23.92it/s, 31 steps of size 1.75e-01. acc. prob=0.79]

warmup:  60%|█████▉    | 448/750 [01:27<00:13, 22.78it/s, 15 steps of size 9.81e-02. acc. prob=0.78]

warmup:  60%|██████    | 450/750 [01:27<00:14, 21.25it/s, 31 steps of size 1.28e-01. acc. prob=0.78]

warmup:  61%|██████    | 454/750 [01:27<00:18, 16.41it/s, 127 steps of size 5.55e-02. acc. prob=0.78]

warmup:  61%|██████    | 455/750 [01:27<00:20, 14.73it/s, 63 steps of size 8.36e-02. acc. prob=0.78]

warmup:  61%|██████    | 457/750 [01:27<00:19, 15.40it/s, 31 steps of size 1.86e-01. acc. prob=0.78]

warmup:  61%|██████▏   | 460/750 [01:27<00:16, 17.11it/s, 31 steps of size 1.94e-01. acc. prob=0.78]

warmup:  62%|██████▏   | 464/750 [01:28<00:15, 19.06it/s, 63 steps of size 7.98e-02. acc. prob=0.78]

warmup:  62%|██████▏   | 465/750 [01:28<00:17, 16.52it/s, 63 steps of size 1.50e-01. acc. prob=0.78]

warmup:  62%|██████▏   | 468/750 [01:28<00:21, 13.14it/s, 127 steps of size 4.72e-02. acc. prob=0.78]

warmup:  63%|██████▎   | 469/750 [01:28<00:27, 10.18it/s, 127 steps of size 8.58e-02. acc. prob=0.78]

warmup:  63%|██████▎   | 470/750 [01:28<00:28,  9.90it/s, 63 steps of size 3.64e-02. acc. prob=0.78]

warmup:  63%|██████▎   | 471/750 [01:29<00:35,  7.94it/s, 127 steps of size 6.61e-02. acc. prob=0.78]

warmup:  63%|██████▎   | 472/750 [01:29<00:34,  8.12it/s, 63 steps of size 1.19e-01. acc. prob=0.78]

warmup:  63%|██████▎   | 475/750 [01:29<00:23, 11.62it/s, 31 steps of size 1.73e-01. acc. prob=0.78]

warmup:  64%|██████▎   | 478/750 [01:29<00:18, 14.34it/s, 31 steps of size 1.85e-01. acc. prob=0.78]

warmup:  64%|██████▍   | 481/750 [01:29<00:18, 14.59it/s, 63 steps of size 1.34e-01. acc. prob=0.78]

warmup:  65%|██████▍   | 484/750 [01:29<00:16, 16.41it/s, 31 steps of size 1.46e-01. acc. prob=0.78]

warmup:  65%|██████▌   | 488/750 [01:30<00:15, 17.43it/s, 63 steps of size 1.38e-01. acc. prob=0.78]

warmup:  65%|██████▌   | 490/750 [01:30<00:14, 17.53it/s, 31 steps of size 2.05e-01. acc. prob=0.78]

warmup:  66%|██████▌   | 494/750 [01:30<00:12, 20.73it/s, 31 steps of size 1.46e-01. acc. prob=0.78]

warmup:  66%|██████▋   | 498/750 [01:30<00:12, 20.19it/s, 63 steps of size 8.28e-02. acc. prob=0.78]

warmup:  67%|██████▋   | 499/750 [01:30<00:14, 17.64it/s, 63 steps of size 1.34e-01. acc. prob=0.78]

sample:  67%|██████▋   | 501/750 [01:30<00:14, 17.68it/s, 31 steps of size 1.48e-01. acc. prob=0.96]

sample:  67%|██████▋   | 503/750 [01:30<00:13, 17.71it/s, 31 steps of size 1.48e-01. acc. prob=0.91]

sample:  67%|██████▋   | 505/750 [01:30<00:13, 17.70it/s, 31 steps of size 1.48e-01. acc. prob=0.80]

sample:  68%|██████▊   | 507/750 [01:31<00:13, 17.71it/s, 31 steps of size 1.48e-01. acc. prob=0.85]

sample:  68%|██████▊   | 509/750 [01:31<00:13, 17.75it/s, 31 steps of size 1.48e-01. acc. prob=0.86]

sample:  68%|██████▊   | 511/750 [01:31<00:13, 17.81it/s, 31 steps of size 1.48e-01. acc. prob=0.87]

sample:  68%|██████▊   | 513/750 [01:31<00:13, 17.84it/s, 31 steps of size 1.48e-01. acc. prob=0.85]

sample:  69%|██████▊   | 515/750 [01:31<00:13, 17.86it/s, 31 steps of size 1.48e-01. acc. prob=0.84]

sample:  69%|██████▉   | 519/750 [01:31<00:11, 20.95it/s, 31 steps of size 1.48e-01. acc. prob=0.78]

sample:  69%|██████▉   | 521/750 [01:31<00:11, 20.09it/s, 31 steps of size 1.48e-01. acc. prob=0.73]

sample:  70%|██████▉   | 523/750 [01:31<00:11, 19.47it/s, 31 steps of size 1.48e-01. acc. prob=0.69]

sample:  70%|███████   | 526/750 [01:32<00:11, 20.34it/s, 31 steps of size 1.48e-01. acc. prob=0.66]

sample:  71%|███████   | 529/750 [01:32<00:10, 20.85it/s, 31 steps of size 1.48e-01. acc. prob=0.67]

sample:  71%|███████   | 531/750 [01:32<00:10, 20.03it/s, 31 steps of size 1.48e-01. acc. prob=0.68]

sample:  71%|███████   | 533/750 [01:32<00:11, 19.42it/s, 31 steps of size 1.48e-01. acc. prob=0.69]

sample:  71%|███████▏  | 535/750 [01:32<00:11, 18.98it/s, 31 steps of size 1.48e-01. acc. prob=0.70]

sample:  72%|███████▏  | 537/750 [01:32<00:11, 18.67it/s, 31 steps of size 1.48e-01. acc. prob=0.72]

sample:  72%|███████▏  | 539/750 [01:32<00:11, 18.44it/s, 31 steps of size 1.48e-01. acc. prob=0.73]

sample:  72%|███████▏  | 541/750 [01:32<00:11, 18.23it/s, 31 steps of size 1.48e-01. acc. prob=0.74]

sample:  72%|███████▏  | 543/750 [01:32<00:11, 17.98it/s, 31 steps of size 1.48e-01. acc. prob=0.75]

sample:  73%|███████▎  | 545/750 [01:33<00:11, 17.92it/s, 31 steps of size 1.48e-01. acc. prob=0.76]

sample:  73%|███████▎  | 547/750 [01:33<00:11, 17.91it/s, 31 steps of size 1.48e-01. acc. prob=0.77]

sample:  73%|███████▎  | 549/750 [01:33<00:11, 17.13it/s, 31 steps of size 1.48e-01. acc. prob=0.78]

sample:  73%|███████▎  | 551/750 [01:33<00:11, 17.26it/s, 31 steps of size 1.48e-01. acc. prob=0.78]

sample:  74%|███████▎  | 553/750 [01:33<00:11, 17.38it/s, 31 steps of size 1.48e-01. acc. prob=0.79]

sample:  74%|███████▍  | 555/750 [01:33<00:11, 17.54it/s, 31 steps of size 1.48e-01. acc. prob=0.80]

sample:  74%|███████▍  | 557/750 [01:33<00:10, 17.62it/s, 31 steps of size 1.48e-01. acc. prob=0.79]

sample:  75%|███████▍  | 559/750 [01:33<00:10, 17.58it/s, 31 steps of size 1.48e-01. acc. prob=0.78]

sample:  75%|███████▌  | 563/750 [01:34<00:08, 20.91it/s, 31 steps of size 1.48e-01. acc. prob=0.76]

sample:  75%|███████▌  | 565/750 [01:34<00:09, 20.07it/s, 31 steps of size 1.48e-01. acc. prob=0.75]

sample:  76%|███████▌  | 567/750 [01:34<00:09, 19.30it/s, 31 steps of size 1.48e-01. acc. prob=0.73]

sample:  76%|███████▌  | 569/750 [01:34<00:09, 18.77it/s, 31 steps of size 1.48e-01. acc. prob=0.74]

sample:  76%|███████▌  | 571/750 [01:34<00:09, 18.50it/s, 31 steps of size 1.48e-01. acc. prob=0.75]

sample:  76%|███████▋  | 573/750 [01:34<00:09, 18.30it/s, 31 steps of size 1.48e-01. acc. prob=0.75]

sample:  77%|███████▋  | 575/750 [01:34<00:09, 17.93it/s, 31 steps of size 1.48e-01. acc. prob=0.75]

sample:  77%|███████▋  | 577/750 [01:34<00:09, 17.80it/s, 31 steps of size 1.48e-01. acc. prob=0.75]

sample:  77%|███████▋  | 579/750 [01:34<00:09, 17.72it/s, 31 steps of size 1.48e-01. acc. prob=0.76]

sample:  77%|███████▋  | 581/750 [01:35<00:09, 17.62it/s, 31 steps of size 1.48e-01. acc. prob=0.76]

sample:  78%|███████▊  | 583/750 [01:35<00:09, 17.63it/s, 31 steps of size 1.48e-01. acc. prob=0.77]

sample:  78%|███████▊  | 585/750 [01:35<00:09, 16.99it/s, 31 steps of size 1.48e-01. acc. prob=0.76]

sample:  78%|███████▊  | 587/750 [01:35<00:09, 17.16it/s, 31 steps of size 1.48e-01. acc. prob=0.76]

sample:  79%|███████▊  | 589/750 [01:35<00:09, 17.32it/s, 31 steps of size 1.48e-01. acc. prob=0.77]

sample:  79%|███████▉  | 591/750 [01:35<00:09, 17.40it/s, 31 steps of size 1.48e-01. acc. prob=0.77]

sample:  79%|███████▉  | 596/750 [01:35<00:06, 24.03it/s, 31 steps of size 1.48e-01. acc. prob=0.74]

sample:  80%|███████▉  | 598/750 [01:35<00:06, 22.22it/s, 31 steps of size 1.48e-01. acc. prob=0.73]

sample:  80%|████████  | 601/750 [01:36<00:06, 22.24it/s, 31 steps of size 1.48e-01. acc. prob=0.72]

sample:  80%|████████  | 603/750 [01:36<00:06, 21.56it/s, 25 steps of size 1.48e-01. acc. prob=0.71]

sample:  81%|████████▏ | 610/750 [01:36<00:04, 34.51it/s, 7 steps of size 1.48e-01. acc. prob=0.68]

sample:  82%|████████▏ | 613/750 [01:36<00:04, 32.46it/s, 19 steps of size 1.48e-01. acc. prob=0.66]

sample:  83%|████████▎ | 622/750 [01:36<00:02, 48.13it/s, 2 steps of size 1.48e-01. acc. prob=0.61]

sample:  83%|████████▎ | 626/750 [01:36<00:02, 45.01it/s, 4 steps of size 1.48e-01. acc. prob=0.60]

sample:  85%|████████▌ | 641/750 [01:36<00:01, 74.61it/s, 4 steps of size 1.48e-01. acc. prob=0.53]

sample:  87%|████████▋ | 655/750 [01:36<00:01, 91.63it/s, 4 steps of size 1.48e-01. acc. prob=0.49]

sample:  89%|████████▉ | 671/750 [01:36<00:00, 109.95it/s, 2 steps of size 1.48e-01. acc. prob=0.45]

sample:  91%|█████████ | 683/750 [01:36<00:00, 108.94it/s, 12 steps of size 1.48e-01. acc. prob=0.42]

sample:  93%|█████████▎| 697/750 [01:37<00:00, 117.76it/s, 7 steps of size 1.48e-01. acc. prob=0.39]

sample:  94%|█████████▎| 702/750 [01:37<00:00, 96.77it/s, 5 steps of size 1.48e-01. acc. prob=0.38]

sample:  95%|█████████▌| 714/750 [01:37<00:00, 100.24it/s, 7 steps of size 1.48e-01. acc. prob=0.36]

sample:  96%|█████████▋| 723/750 [01:37<00:00, 94.80it/s, 6 steps of size 1.48e-01. acc. prob=0.35]

sample:  97%|█████████▋| 727/750 [01:37<00:00, 79.04it/s, 7 steps of size 1.48e-01. acc. prob=0.34]

sample:  98%|█████████▊| 734/750 [01:37<00:00, 74.99it/s, 15 steps of size 1.48e-01. acc. prob=0.33]

sample:  99%|█████████▉| 745/750 [01:37<00:00, 83.39it/s, 7 steps of size 1.48e-01. acc. prob=0.32]

sample: 100%|██████████| 750/750 [01:37<00:00,  7.67it/s, 13 steps of size 1.48e-01. acc. prob=0.32]

  0%|          | 0/750 [00:00<?, ?it/s]

warmup:   5%|▍         | 37/750 [00:10<03:29,  3.40it/s, 511 steps of size 1.31e-03. acc. prob=0.71]

warmup:  10%|▉         | 74/750 [00:21<03:16,  3.44it/s, 7 steps of size 1.70e-03. acc. prob=0.74]

warmup:  15%|█▍        | 111/750 [00:36<03:38,  2.92it/s, 63 steps of size 1.08e-01. acc. prob=0.76]

warmup:  20%|█▉        | 148/750 [00:41<02:40,  3.76it/s, 63 steps of size 8.92e-02. acc. prob=0.77]

warmup:  25%|██▍       | 185/750 [00:46<01:58,  4.76it/s, 15 steps of size 3.01e-01. acc. prob=0.78]

warmup:  30%|██▉       | 222/750 [00:49<01:29,  5.90it/s, 63 steps of size 9.78e-02. acc. prob=0.78]

warmup:  35%|███▍      | 259/750 [00:52<01:08,  7.14it/s, 63 steps of size 8.88e-02. acc. prob=0.78]

warmup:  39%|███▉      | 296/750 [00:54<00:52,  8.69it/s, 15 steps of size 4.07e-01. acc. prob=0.78]

warmup:  44%|████▍     | 333/750 [00:57<00:43,  9.55it/s, 63 steps of size 6.65e-02. acc. prob=0.78]

warmup:  49%|████▉     | 370/750 [01:01<00:39,  9.66it/s, 31 steps of size 6.24e-02. acc. prob=0.78]

warmup:  54%|█████▍    | 407/750 [01:03<00:30, 11.20it/s, 15 steps of size 1.67e-01. acc. prob=0.78]

warmup:  59%|█████▉    | 444/750 [01:06<00:26, 11.49it/s, 159 steps of size 1.19e-01. acc. prob=0.78]

warmup:  64%|██████▍   | 481/750 [01:10<00:24, 11.08it/s, 31 steps of size 1.30e-01. acc. prob=0.78]

sample:  69%|██████▉   | 518/750 [01:13<00:21, 10.69it/s, 63 steps of size 7.90e-02. acc. prob=0.94]

sample:  74%|███████▍  | 555/750 [01:17<00:17, 10.96it/s, 31 steps of size 7.90e-02. acc. prob=0.69]

sample:  79%|███████▉  | 592/750 [01:20<00:15, 10.50it/s, 31 steps of size 7.90e-02. acc. prob=0.69]

sample:  84%|████████▍ | 629/750 [01:24<00:11, 10.32it/s, 47 steps of size 7.90e-02. acc. prob=0.67]

sample:  89%|████████▉ | 666/750 [01:28<00:08, 10.47it/s, 31 steps of size 7.90e-02. acc. prob=0.60]

sample:  94%|█████████▎| 703/750 [01:32<00:04, 10.00it/s, 63 steps of size 7.90e-02. acc. prob=0.67]

sample:  99%|█████████▊| 740/750 [01:36<00:01,  9.61it/s, 63 steps of size 7.90e-02. acc. prob=0.71]

sample: 100%|██████████| 750/750 [01:37<00:00,  7.69it/s, 63 steps of size 7.90e-02. acc. prob=0.72]

  0%|          | 0/750 [00:00<?, ?it/s]

warmup:   5%|▍         | 37/750 [00:11<03:39,  3.24it/s, 511 steps of size 1.22e-03. acc. prob=0.71]

warmup:   5%|▍         | 37/750 [00:22<03:39,  3.24it/s, 511 steps of size 9.92e-04. acc. prob=0.72]

warmup:   7%|▋         | 53/750 [00:22<05:16,  2.20it/s, 511 steps of size 1.13e-03. acc. prob=0.73]

warmup:   7%|▋         | 54/750 [00:23<05:27,  2.13it/s, 511 steps of size 1.88e-03. acc. prob=0.73]

warmup:   7%|▋         | 55/750 [00:23<05:27,  2.12it/s, 255 steps of size 3.06e-03. acc. prob=0.74]

warmup:   8%|▊         | 57/750 [00:24<05:26,  2.13it/s, 511 steps of size 9.85e-04. acc. prob=0.73]

warmup:   8%|▊         | 58/750 [00:25<05:48,  1.99it/s, 511 steps of size 1.58e-03. acc. prob=0.73]

warmup:   8%|▊         | 59/750 [00:26<05:44,  2.00it/s, 255 steps of size 1.52e-03. acc. prob=0.73]

warmup:   8%|▊         | 60/750 [00:26<05:40,  2.03it/s, 255 steps of size 1.20e-03. acc. prob=0.73]

warmup:   8%|▊         | 61/750 [00:27<05:36,  2.05it/s, 255 steps of size 1.45e-03. acc. prob=0.73]

warmup:   8%|▊         | 62/750 [00:27<05:31,  2.07it/s, 255 steps of size 2.20e-03. acc. prob=0.74]

warmup:   9%|▊         | 64/750 [00:29<07:10,  1.59it/s, 1023 steps of size 6.60e-04. acc. prob=0.73]

warmup:   9%|▊         | 65/750 [00:31<09:47,  1.17it/s, 1023 steps of size 1.07e-03. acc. prob=0.73]

warmup:   9%|▉         | 66/750 [00:32<09:54,  1.15it/s, 511 steps of size 1.51e-03. acc. prob=0.74]

warmup:   9%|▉         | 67/750 [00:32<07:53,  1.44it/s, 67 steps of size 1.28e-03. acc. prob=0.74]

warmup:   9%|▉         | 68/750 [00:32<07:08,  1.59it/s, 243 steps of size 1.53e-03. acc. prob=0.74]

warmup:   9%|▉         | 70/750 [00:33<06:24,  1.77it/s, 511 steps of size 5.86e-04. acc. prob=0.73]

warmup:   9%|▉         | 71/750 [00:34<07:17,  1.55it/s, 511 steps of size 8.99e-04. acc. prob=0.74]

warmup:  10%|▉         | 72/750 [00:35<08:02,  1.41it/s, 511 steps of size 1.29e-03. acc. prob=0.74]

warmup:  10%|▉         | 73/750 [00:36<07:59,  1.41it/s, 391 steps of size 1.64e-03. acc. prob=0.74]

warmup:  10%|█         | 75/750 [00:38<09:02,  1.24it/s, 1023 steps of size 5.24e-04. acc. prob=0.74]

warmup:  10%|█         | 76/750 [00:39<11:42,  1.04s/it, 1023 steps of size 7.49e-04. acc. prob=0.74]

warmup:  10%|█         | 77/750 [00:40<11:19,  1.01s/it, 511 steps of size 1.17e-03. acc. prob=0.74]

warmup:  10%|█         | 78/750 [00:41<11:01,  1.02it/s, 511 steps of size 1.22e-03. acc. prob=0.74]

warmup:  11%|█         | 79/750 [00:42<09:23,  1.19it/s, 255 steps of size 1.79e-03. acc. prob=0.75]

warmup:  11%|█         | 80/750 [00:42<07:15,  1.54it/s, 93 steps of size 1.69e-03. acc. prob=0.75]

warmup:  11%|█         | 81/750 [00:42<06:37,  1.68it/s, 255 steps of size 1.57e-03. acc. prob=0.75]

warmup:  11%|█         | 82/750 [00:43<06:10,  1.80it/s, 255 steps of size 1.50e-03. acc. prob=0.75]

warmup:  11%|█         | 83/750 [00:43<05:05,  2.18it/s, 127 steps of size 2.31e-03. acc. prob=0.75]

warmup:  11%|█         | 84/750 [00:43<04:17,  2.59it/s, 119 steps of size 2.32e-03. acc. prob=0.75]

warmup:  11%|█▏        | 85/750 [00:44<04:30,  2.46it/s, 255 steps of size 1.68e-03. acc. prob=0.75]

warmup:  11%|█▏        | 86/750 [00:44<04:39,  2.38it/s, 255 steps of size 2.56e-03. acc. prob=0.75]

warmup:  12%|█▏        | 88/750 [00:46<07:13,  1.53it/s, 1023 steps of size 9.20e-04. acc. prob=0.74]

warmup:  12%|█▏        | 89/750 [00:47<08:02,  1.37it/s, 511 steps of size 1.34e-03. acc. prob=0.75]

warmup:  12%|█▏        | 90/750 [00:47<07:16,  1.51it/s, 255 steps of size 1.93e-03. acc. prob=0.75]

warmup:  12%|█▏        | 91/750 [00:48<06:43,  1.63it/s, 255 steps of size 1.32e-03. acc. prob=0.75]

warmup:  12%|█▏        | 92/750 [00:48<06:18,  1.74it/s, 255 steps of size 2.00e-03. acc. prob=0.75]

warmup:  12%|█▏        | 93/750 [00:49<06:00,  1.82it/s, 255 steps of size 3.02e-03. acc. prob=0.75]

warmup:  13%|█▎        | 94/750 [00:49<04:57,  2.21it/s, 122 steps of size 3.69e-03. acc. prob=0.75]

warmup:  13%|█▎        | 96/750 [00:50<05:04,  2.14it/s, 511 steps of size 1.38e-03. acc. prob=0.75]

warmup:  13%|█▎        | 97/750 [00:50<05:05,  2.14it/s, 255 steps of size 2.05e-03. acc. prob=0.75]

warmup:  13%|█▎        | 98/750 [00:51<05:02,  2.15it/s, 255 steps of size 2.16e-03. acc. prob=0.75]

warmup:  13%|█▎        | 99/750 [00:51<05:01,  2.16it/s, 255 steps of size 2.25e-03. acc. prob=0.75]

warmup:  13%|█▎        | 100/750 [00:51<03:57,  2.74it/s, 63 steps of size 2.15e-03. acc. prob=0.75]

warmup:  13%|█▎        | 101/750 [00:53<08:27,  1.28it/s, 1023 steps of size 3.09e-02. acc. prob=0.76]

warmup:  14%|█▎        | 102/750 [00:54<06:43,  1.61it/s, 127 steps of size 5.46e-02. acc. prob=0.76]

warmup:  14%|█▎        | 103/750 [00:54<05:06,  2.11it/s, 63 steps of size 1.05e-01. acc. prob=0.76]

warmup:  14%|█▍        | 106/750 [00:54<02:50,  3.78it/s, 127 steps of size 4.70e-02. acc. prob=0.76]

warmup:  14%|█▍        | 107/750 [00:54<02:45,  3.89it/s, 127 steps of size 8.69e-02. acc. prob=0.76]

warmup:  14%|█▍        | 108/750 [00:54<02:23,  4.47it/s, 63 steps of size 9.86e-02. acc. prob=0.76]

warmup:  15%|█▍        | 110/750 [00:55<03:30,  3.05it/s, 511 steps of size 1.66e-02. acc. prob=0.76]

warmup:  15%|█▍        | 111/750 [00:56<03:47,  2.81it/s, 255 steps of size 3.11e-02. acc. prob=0.76]

warmup:  15%|█▍        | 112/750 [00:56<03:27,  3.08it/s, 127 steps of size 5.30e-02. acc. prob=0.76]

warmup:  15%|█▌        | 113/750 [00:56<03:11,  3.33it/s, 127 steps of size 8.76e-02. acc. prob=0.76]

warmup:  15%|█▌        | 114/750 [00:56<02:39,  4.00it/s, 63 steps of size 8.10e-02. acc. prob=0.76]

warmup:  15%|█▌        | 115/750 [00:56<02:14,  4.71it/s, 63 steps of size 1.39e-01. acc. prob=0.76]

warmup:  16%|█▌        | 117/750 [00:57<02:22,  4.45it/s, 255 steps of size 2.81e-02. acc. prob=0.76]

warmup:  16%|█▌        | 118/750 [00:58<03:40,  2.86it/s, 383 steps of size 4.95e-02. acc. prob=0.76]

warmup:  16%|█▌        | 119/750 [00:58<03:22,  3.11it/s, 127 steps of size 8.94e-02. acc. prob=0.76]

warmup:  16%|█▌        | 120/750 [00:58<02:48,  3.74it/s, 63 steps of size 1.52e-01. acc. prob=0.77]

warmup:  16%|█▋        | 122/750 [00:58<02:43,  3.85it/s, 255 steps of size 2.58e-02. acc. prob=0.76]

warmup:  16%|█▋        | 123/750 [00:59<03:11,  3.27it/s, 255 steps of size 4.57e-02. acc. prob=0.76]

warmup:  17%|█▋        | 124/750 [00:59<02:58,  3.50it/s, 127 steps of size 7.96e-02. acc. prob=0.77]

warmup:  17%|█▋        | 125/750 [00:59<02:30,  4.17it/s, 63 steps of size 1.09e-01. acc. prob=0.77]

warmup:  17%|█▋        | 127/750 [01:00<02:04,  5.00it/s, 127 steps of size 6.84e-02. acc. prob=0.77]

warmup:  17%|█▋        | 128/750 [01:00<01:52,  5.51it/s, 63 steps of size 9.63e-02. acc. prob=0.77]

warmup:  17%|█▋        | 129/750 [01:00<02:01,  5.11it/s, 127 steps of size 1.54e-01. acc. prob=0.77]

warmup:  17%|█▋        | 131/750 [01:00<01:26,  7.19it/s, 31 steps of size 2.52e-01. acc. prob=0.77]

warmup:  18%|█▊        | 133/750 [01:00<01:23,  7.35it/s, 127 steps of size 7.01e-02. acc. prob=0.77]

warmup:  18%|█▊        | 134/750 [01:00<01:20,  7.61it/s, 63 steps of size 1.18e-01. acc. prob=0.77]

warmup:  18%|█▊        | 136/750 [01:01<01:09,  8.80it/s, 63 steps of size 1.02e-01. acc. prob=0.77]

warmup:  18%|█▊        | 138/750 [01:01<00:57, 10.68it/s, 31 steps of size 1.50e-01. acc. prob=0.77]

warmup:  19%|█▊        | 140/750 [01:01<00:49, 12.25it/s, 31 steps of size 2.14e-01. acc. prob=0.77]

warmup:  19%|█▉        | 142/750 [01:01<00:57, 10.58it/s, 127 steps of size 5.85e-02. acc. prob=0.77]

warmup:  19%|█▉        | 143/750 [01:01<00:59, 10.20it/s, 63 steps of size 9.45e-02. acc. prob=0.77]

warmup:  19%|█▉        | 145/750 [01:01<00:50, 12.01it/s, 31 steps of size 1.61e-01. acc. prob=0.77]

warmup:  20%|█▉        | 147/750 [01:01<00:44, 13.49it/s, 31 steps of size 7.57e-02. acc. prob=0.77]

warmup:  20%|█▉        | 148/750 [01:01<00:49, 12.27it/s, 63 steps of size 1.20e-01. acc. prob=0.77]

warmup:  20%|██        | 150/750 [01:02<01:02,  9.60it/s, 127 steps of size 5.46e-02. acc. prob=0.77]

warmup:  20%|██        | 151/750 [01:02<01:03,  9.38it/s, 63 steps of size 6.14e-01. acc. prob=0.77]

warmup:  21%|██        | 154/750 [01:02<00:44, 13.31it/s, 31 steps of size 1.58e-01. acc. prob=0.77]

warmup:  21%|██        | 156/750 [01:02<01:12,  8.15it/s, 255 steps of size 2.29e-02. acc. prob=0.77]

warmup:  21%|██        | 157/750 [01:03<01:48,  5.48it/s, 255 steps of size 4.05e-02. acc. prob=0.77]

warmup:  21%|██        | 158/750 [01:03<01:53,  5.22it/s, 127 steps of size 7.36e-02. acc. prob=0.77]

warmup:  21%|██        | 159/750 [01:03<01:42,  5.75it/s, 63 steps of size 1.29e-01. acc. prob=0.77]

warmup:  21%|██▏       | 161/750 [01:04<01:33,  6.32it/s, 127 steps of size 6.45e-02. acc. prob=0.77]

warmup:  22%|██▏       | 162/750 [01:04<01:27,  6.74it/s, 63 steps of size 1.10e-01. acc. prob=0.77]

warmup:  22%|██▏       | 165/750 [01:04<01:36,  6.05it/s, 255 steps of size 2.81e-02. acc. prob=0.77]

warmup:  22%|██▏       | 166/750 [01:04<01:43,  5.65it/s, 127 steps of size 5.30e-02. acc. prob=0.77]

warmup:  22%|██▏       | 167/750 [01:05<01:49,  5.32it/s, 127 steps of size 9.32e-02. acc. prob=0.77]

warmup:  22%|██▏       | 168/750 [01:05<01:39,  5.82it/s, 63 steps of size 1.64e-01. acc. prob=0.77]

warmup:  23%|██▎       | 170/750 [01:05<01:34,  6.17it/s, 127 steps of size 6.61e-02. acc. prob=0.77]

warmup:  23%|██▎       | 171/750 [01:05<01:28,  6.56it/s, 63 steps of size 1.02e-01. acc. prob=0.77]

warmup:  23%|██▎       | 172/750 [01:05<01:23,  6.90it/s, 63 steps of size 3.00e-02. acc. prob=0.77]

warmup:  23%|██▎       | 173/750 [01:06<01:37,  5.93it/s, 127 steps of size 5.46e-02. acc. prob=0.77]

warmup:  23%|██▎       | 174/750 [01:06<01:29,  6.42it/s, 63 steps of size 9.76e-02. acc. prob=0.77]

warmup:  23%|██▎       | 176/750 [01:06<01:11,  7.98it/s, 63 steps of size 5.51e-02. acc. prob=0.77]

warmup:  24%|██▎       | 177/750 [01:06<01:10,  8.16it/s, 63 steps of size 9.76e-02. acc. prob=0.77]

warmup:  24%|██▍       | 180/750 [01:06<00:48, 11.86it/s, 31 steps of size 3.06e-01. acc. prob=0.78]

warmup:  24%|██▍       | 182/750 [01:06<00:45, 12.49it/s, 63 steps of size 1.34e-01. acc. prob=0.77]

warmup:  25%|██▍       | 185/750 [01:06<00:35, 16.00it/s, 15 steps of size 2.18e-01. acc. prob=0.78]

warmup:  25%|██▌       | 188/750 [01:07<00:41, 13.54it/s, 127 steps of size 5.99e-02. acc. prob=0.77]

warmup:  25%|██▌       | 189/750 [01:07<00:44, 12.52it/s, 63 steps of size 7.78e-02. acc. prob=0.77]

warmup:  25%|██▌       | 190/750 [01:07<00:48, 11.60it/s, 63 steps of size 5.39e-02. acc. prob=0.77]

warmup:  25%|██▌       | 191/750 [01:07<00:51, 10.86it/s, 63 steps of size 9.03e-02. acc. prob=0.77]

warmup:  26%|██▌       | 193/750 [01:07<00:44, 12.60it/s, 31 steps of size 2.13e-01. acc. prob=0.78]

warmup:  26%|██▌       | 195/750 [01:07<00:39, 13.96it/s, 15 steps of size 1.58e-01. acc. prob=0.78]

warmup:  26%|██▋       | 198/750 [01:07<00:34, 15.84it/s, 31 steps of size 2.26e-01. acc. prob=0.78]

warmup:  27%|██▋       | 200/750 [01:07<00:34, 16.17it/s, 31 steps of size 1.72e-01. acc. prob=0.78]

warmup:  27%|██▋       | 203/750 [01:08<00:28, 19.06it/s, 15 steps of size 5.39e-02. acc. prob=0.77]

warmup:  27%|██▋       | 204/750 [01:08<00:33, 16.10it/s, 63 steps of size 8.41e-02. acc. prob=0.78]

warmup:  27%|██▋       | 206/750 [01:08<00:33, 16.43it/s, 31 steps of size 1.79e-01. acc. prob=0.78]

warmup:  28%|██▊       | 208/750 [01:08<00:32, 16.80it/s, 31 steps of size 1.85e-01. acc. prob=0.78]

warmup:  28%|██▊       | 211/750 [01:08<00:27, 19.64it/s, 15 steps of size 2.38e-01. acc. prob=0.78]

warmup:  29%|██▊       | 214/750 [01:08<00:35, 14.95it/s, 127 steps of size 5.96e-02. acc. prob=0.78]

warmup:  29%|██▊       | 215/750 [01:08<00:39, 13.49it/s, 63 steps of size 9.23e-02. acc. prob=0.78]

warmup:  29%|██▉       | 217/750 [01:09<00:36, 14.51it/s, 31 steps of size 1.64e-01. acc. prob=0.78]

warmup:  29%|██▉       | 220/750 [01:09<00:31, 17.03it/s, 31 steps of size 1.59e-01. acc. prob=0.78]

warmup:  30%|██▉       | 222/750 [01:09<00:34, 15.15it/s, 63 steps of size 1.14e-01. acc. prob=0.78]

warmup:  30%|██▉       | 224/750 [01:09<00:43, 12.10it/s, 127 steps of size 6.48e-02. acc. prob=0.78]

warmup:  30%|███       | 225/750 [01:09<00:46, 11.36it/s, 63 steps of size 8.01e-02. acc. prob=0.78]

warmup:  30%|███       | 226/750 [01:09<00:48, 10.74it/s, 63 steps of size 4.02e-02. acc. prob=0.78]

warmup:  30%|███       | 227/750 [01:10<01:04,  8.05it/s, 127 steps of size 5.54e-02. acc. prob=0.78]

warmup:  30%|███       | 228/750 [01:10<01:03,  8.17it/s, 63 steps of size 6.12e-02. acc. prob=0.78]

warmup:  31%|███       | 229/750 [01:10<01:02,  8.28it/s, 63 steps of size 7.77e-02. acc. prob=0.78]

warmup:  31%|███       | 230/750 [01:10<01:02,  8.37it/s, 63 steps of size 1.60e-02. acc. prob=0.77]

warmup:  31%|███       | 231/750 [01:10<01:51,  4.65it/s, 255 steps of size 2.42e-02. acc. prob=0.77]

warmup:  31%|███       | 232/750 [01:11<02:39,  3.26it/s, 255 steps of size 3.50e-02. acc. prob=0.78]

warmup:  31%|███       | 233/750 [01:11<02:34,  3.34it/s, 127 steps of size 4.65e-02. acc. prob=0.78]

warmup:  31%|███       | 234/750 [01:11<02:10,  3.96it/s, 63 steps of size 1.75e-02. acc. prob=0.77]

warmup:  31%|███▏      | 235/750 [01:12<02:59,  2.87it/s, 255 steps of size 2.60e-02. acc. prob=0.78]

warmup:  31%|███▏      | 236/750 [01:12<02:49,  3.03it/s, 127 steps of size 3.59e-02. acc. prob=0.78]

warmup:  32%|███▏      | 237/750 [01:12<02:41,  3.18it/s, 127 steps of size 5.11e-02. acc. prob=0.78]

warmup:  32%|███▏      | 238/750 [01:13<02:57,  2.88it/s, 191 steps of size 6.74e-02. acc. prob=0.78]

warmup:  32%|███▏      | 239/750 [01:13<02:25,  3.50it/s, 63 steps of size 9.85e-02. acc. prob=0.78]

warmup:  32%|███▏      | 241/750 [01:13<01:42,  4.95it/s, 63 steps of size 8.63e-02. acc. prob=0.78]

warmup:  32%|███▏      | 242/750 [01:13<01:34,  5.36it/s, 63 steps of size 5.27e-02. acc. prob=0.78]

warmup:  32%|███▏      | 243/750 [01:14<01:28,  5.74it/s, 63 steps of size 7.56e-02. acc. prob=0.78]

warmup:  33%|███▎      | 244/750 [01:14<01:23,  6.06it/s, 63 steps of size 1.06e-01. acc. prob=0.78]

warmup:  33%|███▎      | 246/750 [01:14<01:25,  5.87it/s, 127 steps of size 4.67e-02. acc. prob=0.78]

warmup:  33%|███▎      | 247/750 [01:14<01:22,  6.11it/s, 63 steps of size 6.77e-02. acc. prob=0.78]

warmup:  33%|███▎      | 248/750 [01:14<01:19,  6.33it/s, 63 steps of size 9.70e-02. acc. prob=0.78]

warmup:  33%|███▎      | 249/750 [01:14<01:17,  6.43it/s, 63 steps of size 1.39e-01. acc. prob=0.78]

warmup:  33%|███▎      | 251/750 [01:15<00:55,  9.03it/s, 15 steps of size 7.08e-01. acc. prob=0.78]

warmup:  34%|███▍      | 254/750 [01:15<00:45, 10.95it/s, 63 steps of size 8.45e-02. acc. prob=0.78]

warmup:  34%|███▍      | 255/750 [01:15<00:49,  9.95it/s, 63 steps of size 1.19e-01. acc. prob=0.78]

warmup:  34%|███▍      | 257/750 [01:16<01:26,  5.69it/s, 255 steps of size 2.90e-02. acc. prob=0.78]

warmup:  34%|███▍      | 258/750 [01:16<01:37,  5.07it/s, 127 steps of size 5.17e-02. acc. prob=0.78]

warmup:  35%|███▍      | 259/750 [01:16<01:30,  5.41it/s, 63 steps of size 9.55e-02. acc. prob=0.78]

warmup:  35%|███▍      | 261/750 [01:16<01:08,  7.13it/s, 31 steps of size 1.63e-01. acc. prob=0.78]

warmup:  35%|███▌      | 263/750 [01:16<00:59,  8.23it/s, 63 steps of size 1.01e-01. acc. prob=0.78]

warmup:  35%|███▌      | 265/750 [01:16<00:50,  9.63it/s, 31 steps of size 2.99e-01. acc. prob=0.78]

warmup:  36%|███▌      | 267/750 [01:17<00:42, 11.49it/s, 31 steps of size 2.27e-01. acc. prob=0.78]

warmup:  36%|███▌      | 270/750 [01:17<00:37, 12.91it/s, 63 steps of size 9.47e-02. acc. prob=0.78]

warmup:  36%|███▌      | 271/750 [01:17<00:42, 11.38it/s, 63 steps of size 1.54e-01. acc. prob=0.78]

warmup:  36%|███▋      | 273/750 [01:17<00:39, 12.16it/s, 31 steps of size 1.69e-01. acc. prob=0.78]

warmup:  37%|███▋      | 276/750 [01:17<00:29, 15.95it/s, 15 steps of size 4.36e-02. acc. prob=0.78]

warmup:  37%|███▋      | 277/750 [01:17<00:35, 13.27it/s, 63 steps of size 7.83e-02. acc. prob=0.78]

warmup:  37%|███▋      | 278/750 [01:17<00:41, 11.43it/s, 63 steps of size 1.25e-01. acc. prob=0.78]

warmup:  37%|███▋      | 280/750 [01:18<00:38, 12.28it/s, 31 steps of size 2.39e-01. acc. prob=0.78]

warmup:  38%|███▊      | 282/750 [01:18<00:33, 13.95it/s, 31 steps of size 2.44e-01. acc. prob=0.78]

warmup:  38%|███▊      | 284/750 [01:18<00:44, 10.50it/s, 127 steps of size 4.56e-02. acc. prob=0.78]

warmup:  38%|███▊      | 285/750 [01:18<01:00,  7.64it/s, 127 steps of size 7.81e-02. acc. prob=0.78]

warmup:  38%|███▊      | 286/750 [01:18<01:01,  7.49it/s, 63 steps of size 3.68e-02. acc. prob=0.78]

warmup:  38%|███▊      | 287/750 [01:19<01:18,  5.90it/s, 127 steps of size 6.34e-02. acc. prob=0.78]

warmup:  38%|███▊      | 288/750 [01:19<01:15,  6.13it/s, 63 steps of size 1.00e-01. acc. prob=0.78]

warmup:  39%|███▊      | 289/750 [01:19<01:12,  6.32it/s, 63 steps of size 1.35e-01. acc. prob=0.78]

warmup:  39%|███▉      | 292/750 [01:19<01:03,  7.18it/s, 127 steps of size 4.65e-02. acc. prob=0.78]

warmup:  39%|███▉      | 293/750 [01:20<01:17,  5.93it/s, 127 steps of size 7.63e-02. acc. prob=0.78]

warmup:  39%|███▉      | 294/750 [01:20<01:14,  6.14it/s, 63 steps of size 9.23e-02. acc. prob=0.78]

warmup:  39%|███▉      | 295/750 [01:20<01:11,  6.34it/s, 63 steps of size 1.23e-01. acc. prob=0.78]

warmup:  40%|███▉      | 297/750 [01:20<00:55,  8.21it/s, 31 steps of size 8.50e-02. acc. prob=0.78]

warmup:  40%|███▉      | 298/750 [01:20<00:57,  7.91it/s, 63 steps of size 1.23e-01. acc. prob=0.78]

warmup:  40%|████      | 300/750 [01:20<00:52,  8.51it/s, 63 steps of size 8.33e-02. acc. prob=0.78]

warmup:  40%|████      | 301/750 [01:21<00:56,  7.97it/s, 63 steps of size 1.03e-01. acc. prob=0.78]

warmup:  40%|████      | 302/750 [01:21<00:59,  7.56it/s, 63 steps of size 1.53e-01. acc. prob=0.78]

warmup:  41%|████      | 304/750 [01:21<01:03,  7.01it/s, 127 steps of size 4.25e-02. acc. prob=0.78]

warmup:  41%|████      | 305/750 [01:21<01:17,  5.72it/s, 127 steps of size 6.30e-02. acc. prob=0.78]

warmup:  41%|████      | 306/750 [01:21<01:13,  6.00it/s, 63 steps of size 9.88e-02. acc. prob=0.78]

warmup:  41%|████      | 308/750 [01:22<00:56,  7.89it/s, 31 steps of size 4.36e-02. acc. prob=0.78]

warmup:  41%|████      | 309/750 [01:22<01:11,  6.16it/s, 127 steps of size 6.82e-02. acc. prob=0.78]

warmup:  41%|████▏     | 310/750 [01:22<01:09,  6.36it/s, 63 steps of size 1.07e-01. acc. prob=0.78]

warmup:  41%|████▏     | 311/750 [01:22<01:07,  6.52it/s, 63 steps of size 1.48e-01. acc. prob=0.78]

warmup:  42%|████▏     | 313/750 [01:22<00:58,  7.52it/s, 63 steps of size 7.04e-02. acc. prob=0.78]

warmup:  42%|████▏     | 314/750 [01:22<00:58,  7.40it/s, 63 steps of size 1.10e-01. acc. prob=0.78]

warmup:  42%|████▏     | 316/750 [01:23<00:44,  9.67it/s, 18 steps of size 5.70e-02. acc. prob=0.78]

warmup:  42%|████▏     | 317/750 [01:23<00:48,  8.93it/s, 63 steps of size 8.50e-02. acc. prob=0.78]

warmup:  43%|████▎     | 319/750 [01:23<00:47,  9.14it/s, 63 steps of size 1.10e-01. acc. prob=0.78]

warmup:  43%|████▎     | 321/750 [01:23<00:41, 10.44it/s, 31 steps of size 1.93e-01. acc. prob=0.78]

warmup:  43%|████▎     | 323/750 [01:23<00:39, 10.75it/s, 63 steps of size 1.07e-01. acc. prob=0.78]

warmup:  43%|████▎     | 325/750 [01:23<00:36, 11.60it/s, 31 steps of size 1.41e-01. acc. prob=0.78]

warmup:  44%|████▎     | 328/750 [01:24<00:43,  9.69it/s, 127 steps of size 5.93e-02. acc. prob=0.78]

warmup:  44%|████▍     | 329/750 [01:24<00:46,  9.07it/s, 63 steps of size 8.85e-02. acc. prob=0.78]

warmup:  44%|████▍     | 330/750 [01:24<00:49,  8.56it/s, 63 steps of size 9.38e-02. acc. prob=0.78]

warmup:  44%|████▍     | 332/750 [01:24<00:41,  9.96it/s, 31 steps of size 3.59e-02. acc. prob=0.78]

warmup:  44%|████▍     | 333/750 [01:24<00:57,  7.29it/s, 127 steps of size 5.25e-02. acc. prob=0.78]

warmup:  45%|████▍     | 334/750 [01:25<00:57,  7.23it/s, 63 steps of size 7.61e-02. acc. prob=0.78]

warmup:  45%|████▍     | 337/750 [01:25<00:39, 10.49it/s, 23 steps of size 2.93e-02. acc. prob=0.78]

warmup:  45%|████▌     | 338/750 [01:25<00:54,  7.61it/s, 127 steps of size 4.05e-02. acc. prob=0.78]

warmup:  45%|████▌     | 339/750 [01:25<01:07,  6.11it/s, 127 steps of size 5.62e-02. acc. prob=0.78]

warmup:  45%|████▌     | 340/750 [01:26<01:04,  6.31it/s, 63 steps of size 8.35e-02. acc. prob=0.78]

warmup:  46%|████▌     | 342/750 [01:26<00:50,  8.13it/s, 31 steps of size 1.52e-01. acc. prob=0.78]

warmup:  46%|████▌     | 344/750 [01:26<00:42,  9.65it/s, 31 steps of size 2.08e-01. acc. prob=0.78]

warmup:  46%|████▋     | 347/750 [01:26<00:30, 13.32it/s, 31 steps of size 1.39e-01. acc. prob=0.78]

warmup:  47%|████▋     | 349/750 [01:26<00:27, 14.63it/s, 15 steps of size 4.51e-02. acc. prob=0.78]

warmup:  47%|████▋     | 350/750 [01:26<00:42,  9.47it/s, 127 steps of size 6.54e-02. acc. prob=0.78]

warmup:  47%|████▋     | 351/750 [01:26<00:45,  8.85it/s, 63 steps of size 9.14e-02. acc. prob=0.78]

warmup:  47%|████▋     | 353/750 [01:27<00:39,  9.94it/s, 31 steps of size 3.66e-02. acc. prob=0.78]

warmup:  47%|████▋     | 354/750 [01:27<00:54,  7.21it/s, 127 steps of size 5.31e-02. acc. prob=0.78]

warmup:  47%|████▋     | 355/750 [01:27<00:55,  7.17it/s, 63 steps of size 7.61e-02. acc. prob=0.78]

warmup:  47%|████▋     | 356/750 [01:27<00:55,  7.14it/s, 63 steps of size 1.10e-01. acc. prob=0.78]

warmup:  48%|████▊     | 358/750 [01:27<00:43,  9.02it/s, 31 steps of size 1.77e-01. acc. prob=0.78]

warmup:  48%|████▊     | 360/750 [01:27<00:34, 11.22it/s, 15 steps of size 1.01e-01. acc. prob=0.78]

warmup:  48%|████▊     | 362/750 [01:28<00:32, 12.11it/s, 31 steps of size 1.32e-01. acc. prob=0.78]

warmup:  49%|████▊     | 364/750 [01:28<00:29, 13.28it/s, 15 steps of size 2.13e-01. acc. prob=0.78]

warmup:  49%|████▉     | 367/750 [01:28<00:23, 16.34it/s, 15 steps of size 2.64e-01. acc. prob=0.78]

warmup:  49%|████▉     | 370/750 [01:28<00:20, 18.94it/s, 15 steps of size 2.85e-01. acc. prob=0.78]

warmup:  50%|████▉     | 372/750 [01:28<00:20, 18.85it/s, 31 steps of size 1.49e-01. acc. prob=0.78]

warmup:  50%|████▉     | 374/750 [01:28<00:20, 18.57it/s, 15 steps of size 2.02e-01. acc. prob=0.78]

warmup:  50%|█████     | 375/750 [01:28<00:24, 15.54it/s, 47 steps of size 8.05e-02. acc. prob=0.78]

warmup:  50%|█████     | 376/750 [01:28<00:29, 12.53it/s, 63 steps of size 1.09e-01. acc. prob=0.78]

warmup:  50%|█████     | 378/750 [01:29<00:33, 11.17it/s, 63 steps of size 6.69e-02. acc. prob=0.78]

warmup:  51%|█████     | 380/750 [01:29<00:33, 11.06it/s, 63 steps of size 3.32e-02. acc. prob=0.78]

warmup:  51%|█████     | 381/750 [01:29<00:47,  7.79it/s, 127 steps of size 4.45e-02. acc. prob=0.78]

warmup:  51%|█████     | 382/750 [01:29<00:48,  7.63it/s, 63 steps of size 1.39e-02. acc. prob=0.78]

warmup:  51%|█████     | 383/750 [01:30<01:26,  4.23it/s, 255 steps of size 1.95e-02. acc. prob=0.78]

warmup:  51%|█████     | 384/750 [01:30<01:57,  3.11it/s, 255 steps of size 2.65e-02. acc. prob=0.78]

warmup:  51%|█████▏    | 385/750 [01:31<01:53,  3.20it/s, 127 steps of size 2.71e-02. acc. prob=0.78]

warmup:  51%|█████▏    | 386/750 [01:31<01:51,  3.28it/s, 127 steps of size 3.54e-02. acc. prob=0.78]

warmup:  52%|█████▏    | 387/750 [01:31<01:48,  3.33it/s, 127 steps of size 4.81e-02. acc. prob=0.78]

warmup:  52%|█████▏    | 388/750 [01:31<01:32,  3.92it/s, 63 steps of size 1.38e-02. acc. prob=0.78]

warmup:  52%|█████▏    | 389/750 [01:32<02:05,  2.88it/s, 255 steps of size 1.91e-02. acc. prob=0.78]

warmup:  52%|█████▏    | 390/750 [01:33<02:28,  2.42it/s, 255 steps of size 2.65e-02. acc. prob=0.78]

warmup:  52%|█████▏    | 391/750 [01:33<02:14,  2.67it/s, 127 steps of size 3.63e-02. acc. prob=0.78]

warmup:  52%|█████▏    | 393/750 [01:33<02:00,  2.97it/s, 255 steps of size 1.05e-02. acc. prob=0.78]

warmup:  53%|█████▎    | 394/750 [01:35<03:10,  1.87it/s, 511 steps of size 1.41e-02. acc. prob=0.78]

warmup:  53%|█████▎    | 395/750 [01:35<03:12,  1.84it/s, 255 steps of size 1.95e-02. acc. prob=0.78]

warmup:  53%|█████▎    | 396/750 [01:35<02:40,  2.20it/s, 95 steps of size 1.38e-02. acc. prob=0.78]

warmup:  53%|█████▎    | 397/750 [01:36<02:51,  2.06it/s, 255 steps of size 1.83e-02. acc. prob=0.78]

warmup:  53%|█████▎    | 398/750 [01:36<02:59,  1.97it/s, 255 steps of size 2.47e-02. acc. prob=0.78]

warmup:  53%|█████▎    | 399/750 [01:37<02:35,  2.25it/s, 127 steps of size 1.67e-02. acc. prob=0.78]

warmup:  53%|█████▎    | 400/750 [01:37<02:48,  2.08it/s, 255 steps of size 2.26e-02. acc. prob=0.78]

warmup:  53%|█████▎    | 401/750 [01:38<02:27,  2.37it/s, 127 steps of size 2.69e-02. acc. prob=0.78]

warmup:  54%|█████▎    | 402/750 [01:38<02:12,  2.63it/s, 127 steps of size 2.26e-02. acc. prob=0.78]

warmup:  54%|█████▎    | 403/750 [01:39<03:00,  1.92it/s, 383 steps of size 2.91e-02. acc. prob=0.78]

warmup:  54%|█████▍    | 404/750 [01:39<02:35,  2.23it/s, 127 steps of size 3.61e-02. acc. prob=0.78]

warmup:  54%|█████▍    | 405/750 [01:39<02:17,  2.50it/s, 127 steps of size 4.70e-02. acc. prob=0.78]

warmup:  54%|█████▍    | 406/750 [01:39<01:50,  3.11it/s, 63 steps of size 6.22e-02. acc. prob=0.78]

warmup:  54%|█████▍    | 408/750 [01:40<01:46,  3.23it/s, 255 steps of size 2.85e-02. acc. prob=0.78]

warmup:  55%|█████▍    | 409/750 [01:40<01:43,  3.30it/s, 127 steps of size 3.85e-02. acc. prob=0.78]

warmup:  55%|█████▍    | 411/750 [01:41<01:41,  3.34it/s, 255 steps of size 1.61e-02. acc. prob=0.78]

warmup:  55%|█████▍    | 412/750 [01:41<02:01,  2.77it/s, 255 steps of size 2.15e-02. acc. prob=0.78]

warmup:  55%|█████▌    | 413/750 [01:42<01:44,  3.21it/s, 71 steps of size 2.25e-02. acc. prob=0.78]

warmup:  55%|█████▌    | 414/750 [01:42<01:42,  3.29it/s, 127 steps of size 2.84e-02. acc. prob=0.78]

warmup:  55%|█████▌    | 415/750 [01:42<01:39,  3.36it/s, 127 steps of size 3.22e-02. acc. prob=0.78]

warmup:  55%|█████▌    | 416/750 [01:42<01:24,  3.94it/s, 63 steps of size 1.87e-02. acc. prob=0.78]

warmup:  56%|█████▌    | 417/750 [01:43<01:55,  2.89it/s, 255 steps of size 2.34e-02. acc. prob=0.78]

warmup:  56%|█████▌    | 418/750 [01:44<03:10,  1.74it/s, 511 steps of size 3.14e-02. acc. prob=0.78]

warmup:  56%|█████▌    | 419/750 [01:44<02:41,  2.05it/s, 127 steps of size 2.72e-02. acc. prob=0.78]

warmup:  56%|█████▌    | 420/750 [01:45<02:21,  2.34it/s, 127 steps of size 3.42e-02. acc. prob=0.78]

warmup:  56%|█████▌    | 421/750 [01:45<02:06,  2.60it/s, 127 steps of size 4.33e-02. acc. prob=0.78]

warmup:  56%|█████▋    | 422/750 [01:45<01:42,  3.21it/s, 63 steps of size 5.81e-02. acc. prob=0.78]

warmup:  56%|█████▋    | 423/750 [01:45<01:25,  3.83it/s, 63 steps of size 3.29e-02. acc. prob=0.78]

warmup:  57%|█████▋    | 424/750 [01:45<01:27,  3.73it/s, 127 steps of size 4.33e-02. acc. prob=0.78]

warmup:  57%|█████▋    | 425/750 [01:46<01:28,  3.68it/s, 127 steps of size 5.79e-02. acc. prob=0.78]

warmup:  57%|█████▋    | 426/750 [01:46<01:15,  4.30it/s, 63 steps of size 7.69e-02. acc. prob=0.78]

warmup:  57%|█████▋    | 427/750 [01:46<01:06,  4.88it/s, 63 steps of size 8.39e-02. acc. prob=0.78]

warmup:  57%|█████▋    | 429/750 [01:46<00:50,  6.30it/s, 63 steps of size 7.94e-02. acc. prob=0.78]

warmup:  57%|█████▋    | 430/750 [01:46<00:49,  6.49it/s, 63 steps of size 8.28e-02. acc. prob=0.78]

warmup:  57%|█████▋    | 431/750 [01:46<00:47,  6.65it/s, 63 steps of size 1.00e-01. acc. prob=0.78]

warmup:  58%|█████▊    | 433/750 [01:47<00:36,  8.64it/s, 31 steps of size 5.31e-02. acc. prob=0.78]

warmup:  58%|█████▊    | 434/750 [01:47<00:38,  8.16it/s, 63 steps of size 7.09e-02. acc. prob=0.78]

warmup:  58%|█████▊    | 435/750 [01:47<00:40,  7.80it/s, 63 steps of size 8.57e-02. acc. prob=0.78]

warmup:  58%|█████▊    | 437/750 [01:47<00:32,  9.55it/s, 31 steps of size 1.18e-01. acc. prob=0.78]

warmup:  59%|█████▊    | 440/750 [01:47<00:35,  8.82it/s, 127 steps of size 5.35e-02. acc. prob=0.78]

warmup:  59%|█████▉    | 441/750 [01:48<00:36,  8.45it/s, 63 steps of size 6.69e-02. acc. prob=0.78]

warmup:  59%|█████▉    | 442/750 [01:48<00:37,  8.12it/s, 63 steps of size 6.95e-02. acc. prob=0.78]

warmup:  59%|█████▉    | 443/750 [01:48<00:38,  7.88it/s, 63 steps of size 9.18e-02. acc. prob=0.78]

warmup:  59%|█████▉    | 444/750 [01:48<00:39,  7.68it/s, 63 steps of size 7.81e-02. acc. prob=0.78]

warmup:  59%|█████▉    | 445/750 [01:48<00:41,  7.27it/s, 63 steps of size 6.89e-02. acc. prob=0.78]

warmup:  59%|█████▉    | 446/750 [01:48<00:42,  7.16it/s, 63 steps of size 9.10e-02. acc. prob=0.78]

warmup:  60%|█████▉    | 448/750 [01:48<00:33,  9.05it/s, 31 steps of size 1.48e-01. acc. prob=0.79]

warmup:  60%|██████    | 450/750 [01:49<00:28, 10.45it/s, 31 steps of size 2.23e-01. acc. prob=0.79]

warmup:  61%|██████    | 454/750 [01:49<00:20, 14.32it/s, 63 steps of size 6.38e-02. acc. prob=0.78]

warmup:  61%|██████    | 455/750 [01:49<00:23, 12.37it/s, 63 steps of size 9.50e-02. acc. prob=0.78]

warmup:  61%|██████    | 457/750 [01:49<00:22, 12.92it/s, 31 steps of size 2.36e-01. acc. prob=0.78]

warmup:  61%|██████▏   | 460/750 [01:49<00:21, 13.48it/s, 63 steps of size 8.49e-02. acc. prob=0.78]

warmup:  62%|██████▏   | 462/750 [01:49<00:21, 13.69it/s, 31 steps of size 2.60e-01. acc. prob=0.78]

warmup:  62%|██████▏   | 464/750 [01:50<00:21, 13.49it/s, 63 steps of size 1.11e-01. acc. prob=0.78]

warmup:  62%|██████▏   | 466/750 [01:50<00:20, 13.73it/s, 31 steps of size 3.20e-01. acc. prob=0.78]

warmup:  62%|██████▏   | 468/750 [01:50<00:27, 10.26it/s, 127 steps of size 5.35e-02. acc. prob=0.78]

warmup:  63%|██████▎   | 469/750 [01:50<00:36,  7.74it/s, 127 steps of size 8.03e-02. acc. prob=0.78]

warmup:  63%|██████▎   | 470/750 [01:50<00:36,  7.61it/s, 63 steps of size 1.30e-01. acc. prob=0.78]

warmup:  63%|██████▎   | 473/750 [01:51<00:47,  5.86it/s, 255 steps of size 4.20e-02. acc. prob=0.78]

warmup:  63%|██████▎   | 474/750 [01:51<00:53,  5.20it/s, 127 steps of size 7.68e-02. acc. prob=0.78]

warmup:  63%|██████▎   | 476/750 [01:52<00:52,  5.19it/s, 127 steps of size 7.15e-02. acc. prob=0.78]

warmup:  64%|██████▎   | 477/750 [01:52<00:50,  5.41it/s, 63 steps of size 7.84e-02. acc. prob=0.78]

warmup:  64%|██████▎   | 478/750 [01:52<00:49,  5.50it/s, 63 steps of size 1.32e-01. acc. prob=0.78]

warmup:  64%|██████▍   | 480/750 [01:53<01:06,  4.08it/s, 255 steps of size 3.12e-02. acc. prob=0.78]

warmup:  64%|██████▍   | 481/750 [01:53<01:08,  3.95it/s, 127 steps of size 5.47e-02. acc. prob=0.78]

warmup:  64%|██████▍   | 482/750 [01:53<01:00,  4.40it/s, 63 steps of size 9.65e-02. acc. prob=0.78]

warmup:  64%|██████▍   | 483/750 [01:53<00:54,  4.87it/s, 63 steps of size 1.67e-01. acc. prob=0.78]

warmup:  65%|██████▍   | 485/750 [01:54<00:50,  5.29it/s, 127 steps of size 4.26e-02. acc. prob=0.78]

warmup:  65%|██████▍   | 486/750 [01:54<00:55,  4.76it/s, 127 steps of size 7.33e-02. acc. prob=0.78]

warmup:  65%|██████▍   | 487/750 [01:54<00:50,  5.19it/s, 63 steps of size 1.25e-01. acc. prob=0.78]

warmup:  65%|██████▌   | 489/750 [01:54<00:35,  7.46it/s, 15 steps of size 2.73e-01. acc. prob=0.79]

warmup:  65%|██████▌   | 491/750 [01:54<00:29,  8.63it/s, 63 steps of size 8.47e-02. acc. prob=0.78]

warmup:  66%|██████▌   | 492/750 [01:54<00:31,  8.29it/s, 63 steps of size 1.30e-01. acc. prob=0.78]

warmup:  66%|██████▌   | 494/750 [01:55<00:24, 10.58it/s, 15 steps of size 9.93e-02. acc. prob=0.78]

warmup:  66%|██████▌   | 496/750 [01:55<00:21, 11.68it/s, 31 steps of size 2.17e-01. acc. prob=0.79]

warmup:  66%|██████▋   | 498/750 [01:55<00:21, 11.63it/s, 63 steps of size 1.04e-01. acc. prob=0.78]

warmup:  67%|██████▋   | 500/750 [01:55<00:22, 10.89it/s, 63 steps of size 9.22e-02. acc. prob=0.78]

sample:  67%|██████▋   | 501/750 [01:55<00:25,  9.93it/s, 63 steps of size 9.22e-02. acc. prob=0.99]

sample:  67%|██████▋   | 502/750 [01:55<00:26,  9.21it/s, 63 steps of size 9.22e-02. acc. prob=0.66]

sample:  67%|██████▋   | 504/750 [01:56<00:26,  9.36it/s, 63 steps of size 9.22e-02. acc. prob=0.61]

sample:  67%|██████▋   | 505/750 [01:56<00:27,  8.81it/s, 63 steps of size 9.22e-02. acc. prob=0.69]

sample:  68%|██████▊   | 507/750 [01:56<00:23, 10.35it/s, 31 steps of size 9.22e-02. acc. prob=0.78]

sample:  68%|██████▊   | 509/750 [01:56<00:20, 11.55it/s, 31 steps of size 9.22e-02. acc. prob=0.82]

sample:  68%|██████▊   | 511/750 [01:56<00:19, 12.41it/s, 31 steps of size 9.22e-02. acc. prob=0.85]

sample:  68%|██████▊   | 513/750 [01:56<00:20, 11.39it/s, 63 steps of size 9.22e-02. acc. prob=0.87]

sample:  69%|██████▊   | 515/750 [01:56<00:19, 12.23it/s, 31 steps of size 9.22e-02. acc. prob=0.88]

sample:  69%|██████▉   | 516/750 [01:57<00:21, 10.83it/s, 63 steps of size 9.22e-02. acc. prob=0.89]

sample:  69%|██████▉   | 518/750 [01:57<00:22, 10.40it/s, 63 steps of size 9.22e-02. acc. prob=0.90]

sample:  69%|██████▉   | 520/750 [01:57<00:19, 11.52it/s, 31 steps of size 9.22e-02. acc. prob=0.91]

sample:  70%|██████▉   | 522/750 [01:57<00:20, 10.87it/s, 63 steps of size 9.22e-02. acc. prob=0.91]

sample:  70%|██████▉   | 523/750 [01:57<00:22,  9.94it/s, 63 steps of size 9.22e-02. acc. prob=0.90]

sample:  70%|███████   | 525/750 [01:57<00:20, 11.19it/s, 31 steps of size 9.22e-02. acc. prob=0.91]

sample:  70%|███████   | 526/750 [01:58<00:22, 10.10it/s, 63 steps of size 9.22e-02. acc. prob=0.91]

sample:  70%|███████   | 528/750 [01:58<00:19, 11.38it/s, 31 steps of size 9.22e-02. acc. prob=0.92]

sample:  71%|███████   | 529/750 [01:58<00:21, 10.19it/s, 63 steps of size 9.22e-02. acc. prob=0.92]

sample:  71%|███████   | 531/750 [01:58<00:21, 10.00it/s, 63 steps of size 9.22e-02. acc. prob=0.92]

sample:  71%|███████   | 532/750 [01:58<00:23,  9.26it/s, 63 steps of size 9.22e-02. acc. prob=0.92]

sample:  71%|███████   | 533/750 [01:58<00:24,  8.70it/s, 63 steps of size 9.22e-02. acc. prob=0.92]

sample:  71%|███████▏  | 535/750 [01:59<00:29,  7.25it/s, 127 steps of size 9.22e-02. acc. prob=0.90]

sample:  71%|███████▏  | 536/750 [01:59<00:29,  7.24it/s, 63 steps of size 9.22e-02. acc. prob=0.90]

sample:  72%|███████▏  | 537/750 [01:59<00:29,  7.23it/s, 63 steps of size 9.22e-02. acc. prob=0.88]

sample:  72%|███████▏  | 539/750 [01:59<00:23,  9.11it/s, 31 steps of size 9.22e-02. acc. prob=0.85]

sample:  72%|███████▏  | 541/750 [01:59<00:22,  9.34it/s, 63 steps of size 9.22e-02. acc. prob=0.85]

sample:  72%|███████▏  | 543/750 [01:59<00:19, 10.68it/s, 31 steps of size 9.22e-02. acc. prob=0.85]

sample:  73%|███████▎  | 545/750 [02:00<00:17, 11.75it/s, 31 steps of size 9.22e-02. acc. prob=0.84]

sample:  73%|███████▎  | 547/750 [02:00<00:16, 12.53it/s, 31 steps of size 9.22e-02. acc. prob=0.85]

sample:  73%|███████▎  | 548/750 [02:00<00:18, 11.02it/s, 63 steps of size 9.22e-02. acc. prob=0.85]

sample:  73%|███████▎  | 550/750 [02:00<00:18, 10.53it/s, 63 steps of size 9.22e-02. acc. prob=0.85]

sample:  73%|███████▎  | 551/750 [02:00<00:20,  9.64it/s, 63 steps of size 9.22e-02. acc. prob=0.84]

sample:  74%|███████▎  | 552/750 [02:00<00:21,  9.00it/s, 63 steps of size 9.22e-02. acc. prob=0.84]

sample:  74%|███████▎  | 553/750 [02:00<00:23,  8.51it/s, 63 steps of size 9.22e-02. acc. prob=0.83]

sample:  74%|███████▍  | 554/750 [02:01<00:24,  8.16it/s, 63 steps of size 9.22e-02. acc. prob=0.83]

sample:  74%|███████▍  | 555/750 [02:01<00:24,  7.89it/s, 63 steps of size 9.22e-02. acc. prob=0.83]

sample:  74%|███████▍  | 557/750 [02:01<00:21,  8.96it/s, 63 steps of size 9.22e-02. acc. prob=0.83]

sample:  74%|███████▍  | 558/750 [02:01<00:22,  8.47it/s, 63 steps of size 9.22e-02. acc. prob=0.83]

sample:  75%|███████▍  | 560/750 [02:01<00:17, 10.80it/s, 31 steps of size 9.22e-02. acc. prob=0.82]

sample:  75%|███████▍  | 561/750 [02:01<00:19,  9.69it/s, 63 steps of size 9.22e-02. acc. prob=0.81]

sample:  75%|███████▌  | 563/750 [02:01<00:16, 11.16it/s, 31 steps of size 9.22e-02. acc. prob=0.81]

sample:  75%|███████▌  | 565/750 [02:02<00:16, 11.00it/s, 63 steps of size 9.22e-02. acc. prob=0.80]

sample:  75%|███████▌  | 566/750 [02:02<00:18,  9.96it/s, 63 steps of size 9.22e-02. acc. prob=0.80]

sample:  76%|███████▌  | 567/750 [02:02<00:19,  9.19it/s, 63 steps of size 9.22e-02. acc. prob=0.80]

sample:  76%|███████▌  | 568/750 [02:02<00:21,  8.63it/s, 63 steps of size 9.22e-02. acc. prob=0.80]

sample:  76%|███████▌  | 569/750 [02:02<00:22,  8.22it/s, 63 steps of size 9.22e-02. acc. prob=0.81]

sample:  76%|███████▌  | 571/750 [02:02<00:17, 10.11it/s, 31 steps of size 9.22e-02. acc. prob=0.81]

sample:  76%|███████▋  | 573/750 [02:03<00:17,  9.99it/s, 63 steps of size 9.22e-02. acc. prob=0.80]

sample:  77%|███████▋  | 574/750 [02:03<00:19,  9.00it/s, 63 steps of size 9.22e-02. acc. prob=0.80]

sample:  77%|███████▋  | 575/750 [02:03<00:20,  8.40it/s, 63 steps of size 9.22e-02. acc. prob=0.80]

sample:  77%|███████▋  | 576/750 [02:03<00:21,  8.02it/s, 63 steps of size 9.22e-02. acc. prob=0.79]

sample:  77%|███████▋  | 577/750 [02:03<00:22,  7.75it/s, 63 steps of size 9.22e-02. acc. prob=0.79]

sample:  77%|███████▋  | 578/750 [02:03<00:22,  7.56it/s, 63 steps of size 9.22e-02. acc. prob=0.79]

sample:  77%|███████▋  | 579/750 [02:03<00:23,  7.43it/s, 63 steps of size 9.22e-02. acc. prob=0.79]

sample:  77%|███████▋  | 580/750 [02:04<00:23,  7.31it/s, 63 steps of size 9.22e-02. acc. prob=0.79]

sample:  77%|███████▋  | 581/750 [02:04<00:23,  7.22it/s, 63 steps of size 9.22e-02. acc. prob=0.79]

sample:  78%|███████▊  | 582/750 [02:04<00:23,  7.17it/s, 63 steps of size 9.22e-02. acc. prob=0.79]

sample:  78%|███████▊  | 583/750 [02:04<00:23,  7.14it/s, 63 steps of size 9.22e-02. acc. prob=0.79]

sample:  78%|███████▊  | 584/750 [02:04<00:23,  7.12it/s, 63 steps of size 9.22e-02. acc. prob=0.80]

sample:  78%|███████▊  | 585/750 [02:04<00:23,  7.10it/s, 63 steps of size 9.22e-02. acc. prob=0.80]

sample:  78%|███████▊  | 586/750 [02:04<00:23,  7.08it/s, 63 steps of size 9.22e-02. acc. prob=0.80]

sample:  78%|███████▊  | 588/750 [02:05<00:17,  9.10it/s, 31 steps of size 9.22e-02. acc. prob=0.80]

sample:  79%|███████▊  | 590/750 [02:05<00:19,  8.04it/s, 95 steps of size 9.22e-02. acc. prob=0.81]

sample:  79%|███████▉  | 591/750 [02:05<00:20,  7.81it/s, 63 steps of size 9.22e-02. acc. prob=0.81]

sample:  79%|███████▉  | 593/750 [02:05<00:16,  9.45it/s, 31 steps of size 9.22e-02. acc. prob=0.81]

sample:  79%|███████▉  | 595/750 [02:05<00:14, 10.71it/s, 31 steps of size 9.22e-02. acc. prob=0.82]

sample:  80%|███████▉  | 597/750 [02:05<00:14, 10.24it/s, 63 steps of size 9.22e-02. acc. prob=0.82]

sample:  80%|███████▉  | 598/750 [02:06<00:16,  9.43it/s, 63 steps of size 9.22e-02. acc. prob=0.82]

sample:  80%|███████▉  | 599/750 [02:06<00:17,  8.78it/s, 63 steps of size 9.22e-02. acc. prob=0.82]

sample:  80%|████████  | 600/750 [02:06<00:18,  8.32it/s, 63 steps of size 9.22e-02. acc. prob=0.82]

sample:  80%|████████  | 601/750 [02:06<00:18,  7.98it/s, 63 steps of size 9.22e-02. acc. prob=0.82]

sample:  80%|████████  | 602/750 [02:06<00:19,  7.73it/s, 63 steps of size 9.22e-02. acc. prob=0.82]

sample:  81%|████████  | 604/750 [02:06<00:17,  8.49it/s, 63 steps of size 9.22e-02. acc. prob=0.82]

sample:  81%|████████  | 605/750 [02:07<00:17,  8.09it/s, 63 steps of size 9.22e-02. acc. prob=0.82]

sample:  81%|████████  | 606/750 [02:07<00:18,  7.79it/s, 63 steps of size 9.22e-02. acc. prob=0.81]

sample:  81%|████████  | 607/750 [02:07<00:18,  7.64it/s, 63 steps of size 9.22e-02. acc. prob=0.81]

sample:  81%|████████  | 608/750 [02:07<00:17,  8.03it/s, 49 steps of size 9.22e-02. acc. prob=0.81]

sample:  81%|████████▏ | 610/750 [02:07<00:16,  8.59it/s, 63 steps of size 9.22e-02. acc. prob=0.79]

sample:  81%|████████▏ | 611/750 [02:07<00:16,  8.22it/s, 63 steps of size 9.22e-02. acc. prob=0.79]

sample:  82%|████████▏ | 613/750 [02:07<00:15,  8.78it/s, 63 steps of size 9.22e-02. acc. prob=0.77]

sample:  82%|████████▏ | 614/750 [02:08<00:16,  8.39it/s, 63 steps of size 9.22e-02. acc. prob=0.77]

sample:  82%|████████▏ | 615/750 [02:08<00:16,  8.09it/s, 63 steps of size 9.22e-02. acc. prob=0.76]

sample:  82%|████████▏ | 616/750 [02:08<00:17,  7.87it/s, 63 steps of size 9.22e-02. acc. prob=0.75]

sample:  82%|████████▏ | 618/750 [02:08<00:13,  9.66it/s, 43 steps of size 9.22e-02. acc. prob=0.74]

sample:  83%|████████▎ | 619/750 [02:08<00:14,  8.91it/s, 63 steps of size 9.22e-02. acc. prob=0.74]

sample:  83%|████████▎ | 621/750 [02:08<00:12, 10.49it/s, 31 steps of size 9.22e-02. acc. prob=0.74]

sample:  83%|████████▎ | 622/750 [02:08<00:13,  9.47it/s, 63 steps of size 9.22e-02. acc. prob=0.74]

sample:  83%|████████▎ | 624/750 [02:09<00:13,  9.54it/s, 63 steps of size 9.22e-02. acc. prob=0.74]

sample:  83%|████████▎ | 626/750 [02:09<00:14,  8.65it/s, 93 steps of size 9.22e-02. acc. prob=0.73]

sample:  84%|████████▎ | 627/750 [02:09<00:14,  8.34it/s, 63 steps of size 9.22e-02. acc. prob=0.73]

sample:  84%|████████▍ | 629/750 [02:09<00:12,  9.89it/s, 31 steps of size 9.22e-02. acc. prob=0.73]

sample:  84%|████████▍ | 631/750 [02:09<00:10, 11.15it/s, 31 steps of size 9.22e-02. acc. prob=0.73]

sample:  84%|████████▍ | 632/750 [02:09<00:11, 10.04it/s, 63 steps of size 9.22e-02. acc. prob=0.73]

sample:  84%|████████▍ | 633/750 [02:10<00:12,  9.16it/s, 63 steps of size 9.22e-02. acc. prob=0.73]

sample:  85%|████████▍ | 634/750 [02:10<00:14,  8.01it/s, 63 steps of size 9.22e-02. acc. prob=0.73]

sample:  85%|████████▍ | 636/750 [02:10<00:12,  8.98it/s, 40 steps of size 9.22e-02. acc. prob=0.72]

sample:  85%|████████▌ | 639/750 [02:10<00:08, 13.31it/s, 3 steps of size 9.22e-02. acc. prob=0.71]

sample:  86%|████████▌ | 642/750 [02:10<00:06, 16.57it/s, 11 steps of size 9.22e-02. acc. prob=0.69]

sample:  86%|████████▌ | 644/750 [02:10<00:07, 14.46it/s, 63 steps of size 9.22e-02. acc. prob=0.69]

sample:  86%|████████▌ | 646/750 [02:11<00:07, 13.72it/s, 63 steps of size 9.22e-02. acc. prob=0.68]

sample:  86%|████████▋ | 647/750 [02:11<00:08, 11.80it/s, 63 steps of size 9.22e-02. acc. prob=0.67]

sample:  86%|████████▋ | 648/750 [02:11<00:09, 10.44it/s, 63 steps of size 9.22e-02. acc. prob=0.67]

sample:  87%|████████▋ | 650/750 [02:11<00:08, 11.98it/s, 31 steps of size 9.22e-02. acc. prob=0.67]

sample:  87%|████████▋ | 653/750 [02:11<00:06, 14.98it/s, 31 steps of size 9.22e-02. acc. prob=0.66]

sample:  87%|████████▋ | 655/750 [02:11<00:06, 15.17it/s, 31 steps of size 9.22e-02. acc. prob=0.66]

sample:  88%|████████▊ | 657/750 [02:11<00:05, 15.71it/s, 21 steps of size 9.22e-02. acc. prob=0.66]

sample:  88%|████████▊ | 659/750 [02:11<00:06, 14.15it/s, 39 steps of size 9.22e-02. acc. prob=0.65]

sample:  88%|████████▊ | 660/750 [02:12<00:07, 11.71it/s, 63 steps of size 9.22e-02. acc. prob=0.65]

sample:  88%|████████▊ | 661/750 [02:12<00:08,  9.94it/s, 63 steps of size 9.22e-02. acc. prob=0.65]

sample:  89%|████████▉ | 666/750 [02:12<00:04, 18.75it/s, 3 steps of size 9.22e-02. acc. prob=0.63]

sample:  89%|████████▉ | 671/750 [02:12<00:03, 25.73it/s, 9 steps of size 9.22e-02. acc. prob=0.62]

sample:  90%|████████▉ | 673/750 [02:12<00:03, 19.74it/s, 63 steps of size 9.22e-02. acc. prob=0.61]

sample:  90%|█████████ | 675/750 [02:12<00:04, 18.32it/s, 36 steps of size 9.22e-02. acc. prob=0.61]

sample:  90%|█████████ | 677/750 [02:12<00:03, 18.51it/s, 23 steps of size 9.22e-02. acc. prob=0.60]

sample:  91%|█████████▏| 685/750 [02:13<00:01, 33.48it/s, 3 steps of size 9.22e-02. acc. prob=0.58]

sample:  93%|█████████▎| 698/750 [02:13<00:00, 58.46it/s, 3 steps of size 9.22e-02. acc. prob=0.54]

sample:  95%|█████████▍| 712/750 [02:13<00:00, 79.28it/s, 3 steps of size 9.22e-02. acc. prob=0.51]

sample:  97%|█████████▋| 727/750 [02:13<00:00, 96.93it/s, 3 steps of size 9.22e-02. acc. prob=0.47]

sample:  99%|█████████▉| 742/750 [02:13<00:00, 109.47it/s, 4 steps of size 9.22e-02. acc. prob=0.45]

sample: 100%|██████████| 750/750 [02:13<00:00,  5.62it/s, 5 steps of size 9.22e-02. acc. prob=0.44]

  0%|          | 0/750 [00:00<?, ?it/s]

warmup:   5%|▍         | 37/750 [00:16<05:11,  2.29it/s, 255 steps of size 2.64e-03. acc. prob=0.71]

warmup:   5%|▍         | 37/750 [00:28<05:11,  2.29it/s, 511 steps of size 1.47e-03. acc. prob=0.73]

warmup:   7%|▋         | 51/750 [00:28<06:58,  1.67it/s, 255 steps of size 2.42e-03. acc. prob=0.73]

warmup:   7%|▋         | 53/750 [00:31<07:25,  1.56it/s, 1023 steps of size 6.90e-04. acc. prob=0.72]

warmup:   7%|▋         | 54/750 [00:33<08:20,  1.39it/s, 1023 steps of size 1.02e-03. acc. prob=0.73]

warmup:   7%|▋         | 55/750 [00:34<08:36,  1.35it/s, 511 steps of size 1.59e-03. acc. prob=0.73]

warmup:   7%|▋         | 56/750 [00:35<08:27,  1.37it/s, 255 steps of size 2.17e-03. acc. prob=0.73]

warmup:   8%|▊         | 58/750 [00:37<09:35,  1.20it/s, 1023 steps of size 8.53e-04. acc. prob=0.73]

warmup:   8%|▊         | 59/750 [00:39<10:06,  1.14it/s, 511 steps of size 1.30e-03. acc. prob=0.73]

warmup:   8%|▊         | 60/750 [00:39<09:37,  1.19it/s, 255 steps of size 9.34e-04. acc. prob=0.73]

warmup:   8%|▊         | 61/750 [00:40<09:03,  1.27it/s, 255 steps of size 1.31e-03. acc. prob=0.73]

warmup:   8%|▊         | 62/750 [00:40<08:31,  1.34it/s, 255 steps of size 1.55e-03. acc. prob=0.74]

warmup:   8%|▊         | 63/750 [00:41<08:05,  1.41it/s, 255 steps of size 3.70e-04. acc. prob=0.73]

warmup:   9%|▊         | 64/750 [00:43<12:17,  1.07s/it, 1023 steps of size 6.04e-04. acc. prob=0.73]

warmup:   9%|▊         | 65/750 [00:45<15:50,  1.39s/it, 1023 steps of size 8.58e-04. acc. prob=0.73]

warmup:   9%|▉         | 66/750 [00:46<15:01,  1.32s/it, 511 steps of size 1.38e-03. acc. prob=0.74]

warmup:   9%|▉         | 67/750 [00:47<12:38,  1.11s/it, 255 steps of size 1.55e-03. acc. prob=0.74]

warmup:   9%|▉         | 68/750 [00:48<10:52,  1.05it/s, 255 steps of size 1.48e-03. acc. prob=0.74]

warmup:   9%|▉         | 69/750 [00:48<09:35,  1.18it/s, 255 steps of size 1.87e-03. acc. prob=0.74]

warmup:   9%|▉         | 70/750 [00:49<08:44,  1.30it/s, 255 steps of size 7.62e-04. acc. prob=0.74]

warmup:   9%|▉         | 71/750 [00:50<09:58,  1.13it/s, 511 steps of size 1.19e-03. acc. prob=0.74]

warmup:  10%|▉         | 72/750 [00:51<10:46,  1.05it/s, 511 steps of size 1.87e-03. acc. prob=0.74]

warmup:  10%|▉         | 73/750 [00:52<09:28,  1.19it/s, 255 steps of size 2.87e-03. acc. prob=0.75]

warmup:  10%|▉         | 74/750 [00:52<07:34,  1.49it/s, 127 steps of size 1.07e-03. acc. prob=0.74]

warmup:  10%|█         | 75/750 [00:53<08:56,  1.26it/s, 492 steps of size 1.37e-03. acc. prob=0.74]

warmup:  10%|█         | 76/750 [00:54<08:14,  1.36it/s, 255 steps of size 2.09e-03. acc. prob=0.75]

warmup:  10%|█         | 77/750 [00:54<07:39,  1.46it/s, 255 steps of size 3.11e-03. acc. prob=0.75]

warmup:  10%|█         | 78/750 [00:54<05:49,  1.92it/s, 63 steps of size 6.58e-04. acc. prob=0.74]

warmup:  11%|█         | 79/750 [00:57<11:34,  1.04s/it, 1023 steps of size 9.86e-04. acc. prob=0.74]

warmup:  11%|█         | 80/750 [00:58<11:53,  1.06s/it, 511 steps of size 1.53e-03. acc. prob=0.74]

warmup:  11%|█         | 81/750 [00:59<12:11,  1.09s/it, 511 steps of size 1.88e-03. acc. prob=0.75]

warmup:  11%|█         | 82/750 [00:59<09:26,  1.18it/s, 127 steps of size 4.36e-04. acc. prob=0.74]

warmup:  11%|█         | 83/750 [01:01<14:07,  1.27s/it, 1023 steps of size 6.58e-04. acc. prob=0.74]

warmup:  11%|█         | 84/750 [01:04<17:27,  1.57s/it, 1023 steps of size 9.97e-04. acc. prob=0.74]

warmup:  11%|█▏        | 85/750 [01:05<15:55,  1.44s/it, 511 steps of size 1.51e-03. acc. prob=0.75]

warmup:  11%|█▏        | 86/750 [01:05<12:59,  1.17s/it, 255 steps of size 1.37e-03. acc. prob=0.75]

warmup:  12%|█▏        | 87/750 [01:06<10:57,  1.01it/s, 255 steps of size 1.75e-03. acc. prob=0.75]

warmup:  12%|█▏        | 88/750 [01:06<09:32,  1.16it/s, 255 steps of size 1.65e-03. acc. prob=0.75]

warmup:  12%|█▏        | 89/750 [01:07<08:37,  1.28it/s, 255 steps of size 2.25e-03. acc. prob=0.75]

warmup:  12%|█▏        | 91/750 [01:09<10:26,  1.05it/s, 1023 steps of size 8.25e-04. acc. prob=0.75]

warmup:  12%|█▏        | 92/750 [01:12<13:55,  1.27s/it, 1023 steps of size 1.24e-03. acc. prob=0.75]

warmup:  12%|█▏        | 93/750 [01:12<11:52,  1.08s/it, 255 steps of size 1.78e-03. acc. prob=0.75]

warmup:  13%|█▎        | 94/750 [01:13<10:17,  1.06it/s, 255 steps of size 1.66e-03. acc. prob=0.75]

warmup:  13%|█▎        | 95/750 [01:13<09:20,  1.17it/s, 255 steps of size 2.00e-03. acc. prob=0.75]

warmup:  13%|█▎        | 96/750 [01:14<08:29,  1.28it/s, 255 steps of size 7.52e-04. acc. prob=0.75]

warmup:  13%|█▎        | 97/750 [01:16<13:40,  1.26s/it, 1023 steps of size 1.13e-03. acc. prob=0.75]

warmup:  13%|█▎        | 98/750 [01:18<13:36,  1.25s/it, 511 steps of size 4.53e-04. acc. prob=0.74]

warmup:  13%|█▎        | 99/750 [01:20<16:39,  1.53s/it, 1023 steps of size 6.78e-04. acc. prob=0.75]

warmup:  13%|█▎        | 100/750 [01:22<18:47,  1.73s/it, 1023 steps of size 8.65e-04. acc. prob=0.75]

warmup:  13%|█▎        | 101/750 [01:24<20:18,  1.88s/it, 1023 steps of size 1.24e-02. acc. prob=0.75]

warmup:  14%|█▎        | 102/750 [01:25<17:46,  1.65s/it, 511 steps of size 2.20e-02. acc. prob=0.75]

warmup:  14%|█▎        | 103/750 [01:26<14:14,  1.32s/it, 255 steps of size 4.16e-02. acc. prob=0.76]

warmup:  14%|█▍        | 104/750 [01:26<10:53,  1.01s/it, 127 steps of size 7.85e-02. acc. prob=0.76]

warmup:  14%|█▍        | 105/750 [01:26<08:03,  1.33it/s, 63 steps of size 2.27e-02. acc. prob=0.75]

warmup:  14%|█▍        | 106/750 [01:27<07:25,  1.45it/s, 255 steps of size 4.28e-02. acc. prob=0.76]

warmup:  14%|█▍        | 107/750 [01:27<06:05,  1.76it/s, 127 steps of size 8.12e-02. acc. prob=0.76]

warmup:  14%|█▍        | 108/750 [01:27<04:42,  2.27it/s, 63 steps of size 8.80e-02. acc. prob=0.76]

warmup:  15%|█▍        | 109/750 [01:27<03:44,  2.85it/s, 63 steps of size 1.03e-01. acc. prob=0.76]

warmup:  15%|█▍        | 110/750 [01:28<03:04,  3.48it/s, 63 steps of size 2.06e-02. acc. prob=0.76]

warmup:  15%|█▍        | 111/750 [01:28<03:56,  2.71it/s, 255 steps of size 3.78e-02. acc. prob=0.76]

warmup:  15%|█▍        | 112/750 [01:28<03:38,  2.92it/s, 127 steps of size 6.99e-02. acc. prob=0.76]

warmup:  15%|█▌        | 113/750 [01:29<02:59,  3.55it/s, 63 steps of size 5.30e-02. acc. prob=0.76]

warmup:  15%|█▌        | 114/750 [01:29<02:31,  4.18it/s, 63 steps of size 6.46e-02. acc. prob=0.76]

warmup:  15%|█▌        | 115/750 [01:29<02:39,  3.99it/s, 127 steps of size 1.11e-01. acc. prob=0.76]

warmup:  16%|█▌        | 117/750 [01:30<02:50,  3.72it/s, 255 steps of size 3.03e-02. acc. prob=0.76]

warmup:  16%|█▌        | 118/750 [01:30<03:35,  2.94it/s, 255 steps of size 5.10e-02. acc. prob=0.76]

warmup:  16%|█▌        | 119/750 [01:30<03:24,  3.09it/s, 127 steps of size 7.51e-02. acc. prob=0.76]

warmup:  16%|█▌        | 120/750 [01:30<02:52,  3.66it/s, 63 steps of size 6.81e-02. acc. prob=0.76]

warmup:  16%|█▌        | 121/750 [01:31<02:27,  4.25it/s, 63 steps of size 1.11e-01. acc. prob=0.76]

warmup:  16%|█▋        | 123/750 [01:31<01:50,  5.68it/s, 63 steps of size 9.00e-02. acc. prob=0.76]

warmup:  17%|█▋        | 124/750 [01:31<01:44,  6.00it/s, 63 steps of size 5.43e-02. acc. prob=0.76]

warmup:  17%|█▋        | 125/750 [01:31<02:02,  5.12it/s, 127 steps of size 7.96e-02. acc. prob=0.77]

warmup:  17%|█▋        | 126/750 [01:31<01:52,  5.56it/s, 63 steps of size 1.30e-01. acc. prob=0.77]

warmup:  17%|█▋        | 128/750 [01:32<01:30,  6.84it/s, 63 steps of size 1.09e-01. acc. prob=0.77]

warmup:  17%|█▋        | 129/750 [01:32<01:29,  6.94it/s, 63 steps of size 1.68e-01. acc. prob=0.77]

warmup:  18%|█▊        | 132/750 [01:33<02:19,  4.42it/s, 383 steps of size 5.27e-02. acc. prob=0.77]

warmup:  18%|█▊        | 133/750 [01:33<02:25,  4.24it/s, 127 steps of size 8.58e-02. acc. prob=0.77]

warmup:  18%|█▊        | 134/750 [01:33<02:12,  4.66it/s, 63 steps of size 1.42e-01. acc. prob=0.77]

warmup:  18%|█▊        | 136/750 [01:33<01:33,  6.55it/s, 15 steps of size 1.63e-01. acc. prob=0.77]

warmup:  18%|█▊        | 138/750 [01:33<01:11,  8.53it/s, 15 steps of size 9.83e-02. acc. prob=0.77]

warmup:  19%|█▊        | 139/750 [01:33<01:15,  8.09it/s, 63 steps of size 1.60e-01. acc. prob=0.77]

warmup:  19%|█▉        | 141/750 [01:34<01:00, 10.09it/s, 15 steps of size 1.18e-01. acc. prob=0.77]

warmup:  19%|█▉        | 142/750 [01:34<01:05,  9.26it/s, 63 steps of size 1.91e-01. acc. prob=0.77]

warmup:  19%|█▉        | 144/750 [01:34<00:53, 11.38it/s, 15 steps of size 5.12e-02. acc. prob=0.77]

warmup:  19%|█▉        | 145/750 [01:34<01:24,  7.13it/s, 127 steps of size 8.21e-02. acc. prob=0.77]

warmup:  19%|█▉        | 146/750 [01:34<01:26,  6.99it/s, 63 steps of size 1.04e-01. acc. prob=0.77]

warmup:  20%|█▉        | 147/750 [01:34<01:28,  6.85it/s, 63 steps of size 4.78e-02. acc. prob=0.77]

warmup:  20%|█▉        | 148/750 [01:35<01:55,  5.22it/s, 127 steps of size 7.51e-02. acc. prob=0.77]

warmup:  20%|█▉        | 149/750 [01:35<01:48,  5.53it/s, 63 steps of size 1.19e-01. acc. prob=0.77]

warmup:  20%|██        | 151/750 [01:35<01:23,  7.18it/s, 31 steps of size 2.21e+00. acc. prob=0.78]

warmup:  21%|██        | 154/750 [01:35<01:16,  7.74it/s, 127 steps of size 5.19e-02. acc. prob=0.77]

warmup:  21%|██        | 155/750 [01:36<01:18,  7.60it/s, 63 steps of size 8.03e-02. acc. prob=0.77]

warmup:  21%|██        | 156/750 [01:36<01:23,  7.13it/s, 63 steps of size 1.31e-01. acc. prob=0.77]

warmup:  21%|██        | 158/750 [01:36<01:07,  8.71it/s, 31 steps of size 2.29e-02. acc. prob=0.77]

warmup:  21%|██        | 159/750 [01:36<01:31,  6.43it/s, 127 steps of size 4.21e-02. acc. prob=0.77]

warmup:  21%|██▏       | 160/750 [01:36<01:32,  6.39it/s, 63 steps of size 7.76e-02. acc. prob=0.77]

warmup:  21%|██▏       | 161/750 [01:36<01:31,  6.47it/s, 63 steps of size 1.38e-01. acc. prob=0.77]

warmup:  22%|██▏       | 163/750 [01:37<01:05,  8.94it/s, 15 steps of size 3.82e-01. acc. prob=0.77]

warmup:  22%|██▏       | 166/750 [01:37<00:45, 12.95it/s, 15 steps of size 2.46e-01. acc. prob=0.77]

warmup:  22%|██▏       | 168/750 [01:37<00:57, 10.06it/s, 127 steps of size 5.07e-02. acc. prob=0.77]

warmup:  23%|██▎       | 169/750 [01:37<01:01,  9.43it/s, 63 steps of size 9.40e-02. acc. prob=0.77]

warmup:  23%|██▎       | 171/750 [01:38<02:44,  3.52it/s, 511 steps of size 1.88e-02. acc. prob=0.77]

warmup:  23%|██▎       | 172/750 [01:39<03:26,  2.80it/s, 255 steps of size 3.36e-02. acc. prob=0.77]

warmup:  23%|██▎       | 173/750 [01:39<03:16,  2.93it/s, 127 steps of size 6.08e-02. acc. prob=0.77]

warmup:  23%|██▎       | 175/750 [01:40<03:02,  3.15it/s, 255 steps of size 2.04e-02. acc. prob=0.77]

warmup:  24%|██▎       | 177/750 [01:42<05:51,  1.63it/s, 1023 steps of size 4.08e-03. acc. prob=0.77]

warmup:  24%|██▎       | 178/750 [01:45<09:06,  1.05it/s, 1023 steps of size 7.20e-03. acc. prob=0.77]

warmup:  24%|██▍       | 179/750 [01:46<09:27,  1.01it/s, 511 steps of size 1.10e-02. acc. prob=0.77]

warmup:  24%|██▍       | 180/750 [01:47<09:41,  1.02s/it, 511 steps of size 1.36e-02. acc. prob=0.77]

warmup:  24%|██▍       | 181/750 [01:47<08:47,  1.08it/s, 255 steps of size 7.63e-03. acc. prob=0.77]

warmup:  24%|██▍       | 182/750 [01:49<09:21,  1.01it/s, 511 steps of size 1.33e-02. acc. prob=0.77]

warmup:  24%|██▍       | 183/750 [01:49<08:24,  1.12it/s, 255 steps of size 2.07e-02. acc. prob=0.77]

warmup:  25%|██▍       | 184/750 [01:50<07:31,  1.25it/s, 255 steps of size 2.87e-02. acc. prob=0.77]

warmup:  25%|██▍       | 185/750 [01:50<06:05,  1.54it/s, 127 steps of size 8.75e-03. acc. prob=0.77]

warmup:  25%|██▍       | 186/750 [01:51<07:23,  1.27it/s, 511 steps of size 1.54e-02. acc. prob=0.77]

warmup:  25%|██▍       | 187/750 [01:52<06:43,  1.40it/s, 255 steps of size 2.66e-02. acc. prob=0.77]

warmup:  25%|██▌       | 188/750 [01:52<05:30,  1.70it/s, 127 steps of size 2.82e-02. acc. prob=0.77]

warmup:  25%|██▌       | 189/750 [01:52<04:38,  2.01it/s, 127 steps of size 4.11e-02. acc. prob=0.77]

warmup:  25%|██▌       | 190/750 [01:53<04:01,  2.32it/s, 127 steps of size 5.12e-02. acc. prob=0.77]

warmup:  25%|██▌       | 191/750 [01:53<03:12,  2.90it/s, 63 steps of size 7.93e-02. acc. prob=0.77]

warmup:  26%|██▌       | 193/750 [01:53<03:03,  3.04it/s, 255 steps of size 2.10e-02. acc. prob=0.77]

warmup:  26%|██▌       | 194/750 [01:54<03:45,  2.46it/s, 255 steps of size 3.39e-02. acc. prob=0.77]

warmup:  26%|██▌       | 195/750 [01:54<03:28,  2.67it/s, 127 steps of size 1.20e-02. acc. prob=0.77]

warmup:  26%|██▌       | 196/750 [01:55<04:00,  2.30it/s, 255 steps of size 1.99e-02. acc. prob=0.77]

warmup:  26%|██▋       | 197/750 [01:55<04:18,  2.14it/s, 255 steps of size 2.48e-02. acc. prob=0.77]

warmup:  26%|██▋       | 198/750 [01:56<03:47,  2.42it/s, 127 steps of size 9.31e-03. acc. prob=0.77]

warmup:  27%|██▋       | 199/750 [01:57<05:36,  1.64it/s, 511 steps of size 1.45e-02. acc. prob=0.77]

warmup:  27%|██▋       | 201/750 [01:59<07:41,  1.19it/s, 1023 steps of size 4.47e-03. acc. prob=0.77]

warmup:  27%|██▋       | 202/750 [02:01<10:51,  1.19s/it, 1023 steps of size 6.86e-03. acc. prob=0.77]

warmup:  27%|██▋       | 203/750 [02:02<10:41,  1.17s/it, 511 steps of size 1.12e-02. acc. prob=0.77]

warmup:  27%|██▋       | 204/750 [02:03<09:08,  1.01s/it, 255 steps of size 8.67e-03. acc. prob=0.77]

warmup:  27%|██▋       | 205/750 [02:04<09:35,  1.06s/it, 511 steps of size 1.25e-02. acc. prob=0.77]

warmup:  27%|██▋       | 206/750 [02:05<08:21,  1.09it/s, 255 steps of size 2.02e-02. acc. prob=0.77]

warmup:  28%|██▊       | 207/750 [02:05<07:24,  1.22it/s, 255 steps of size 3.84e-03. acc. prob=0.77]

warmup:  28%|██▊       | 208/750 [02:08<11:13,  1.24s/it, 1023 steps of size 6.23e-03. acc. prob=0.77]

warmup:  28%|██▊       | 209/750 [02:09<10:52,  1.21s/it, 511 steps of size 9.98e-03. acc. prob=0.77]

warmup:  28%|██▊       | 210/750 [02:10<10:32,  1.17s/it, 511 steps of size 1.55e-02. acc. prob=0.77]

warmup:  28%|██▊       | 211/750 [02:10<08:51,  1.01it/s, 255 steps of size 2.46e-02. acc. prob=0.77]

warmup:  28%|██▊       | 212/750 [02:11<06:55,  1.29it/s, 127 steps of size 1.29e-02. acc. prob=0.77]

warmup:  28%|██▊       | 213/750 [02:11<06:18,  1.42it/s, 255 steps of size 1.89e-02. acc. prob=0.77]

warmup:  29%|██▊       | 214/750 [02:12<05:52,  1.52it/s, 255 steps of size 2.97e-02. acc. prob=0.77]

warmup:  29%|██▊       | 215/750 [02:12<04:51,  1.83it/s, 127 steps of size 4.23e-02. acc. prob=0.77]

warmup:  29%|██▉       | 216/750 [02:12<04:08,  2.15it/s, 127 steps of size 6.35e-02. acc. prob=0.77]

warmup:  29%|██▉       | 217/750 [02:12<03:15,  2.72it/s, 63 steps of size 9.89e-02. acc. prob=0.78]

warmup:  29%|██▉       | 219/750 [02:13<02:59,  2.96it/s, 255 steps of size 2.86e-02. acc. prob=0.77]

warmup:  29%|██▉       | 220/750 [02:13<02:52,  3.07it/s, 127 steps of size 3.99e-02. acc. prob=0.77]

warmup:  29%|██▉       | 221/750 [02:14<02:45,  3.20it/s, 127 steps of size 6.04e-02. acc. prob=0.78]

warmup:  30%|██▉       | 222/750 [02:14<02:19,  3.78it/s, 63 steps of size 9.11e-02. acc. prob=0.78]

warmup:  30%|██▉       | 223/750 [02:14<02:21,  3.74it/s, 127 steps of size 2.89e-02. acc. prob=0.77]

warmup:  30%|██▉       | 224/750 [02:14<02:25,  3.61it/s, 127 steps of size 4.19e-02. acc. prob=0.77]

warmup:  30%|███       | 225/750 [02:15<02:27,  3.57it/s, 127 steps of size 5.92e-02. acc. prob=0.78]

warmup:  30%|███       | 226/750 [02:15<02:07,  4.11it/s, 63 steps of size 7.72e-02. acc. prob=0.78]

warmup:  30%|███       | 227/750 [02:15<01:51,  4.67it/s, 63 steps of size 1.13e-01. acc. prob=0.78]

warmup:  30%|███       | 228/750 [02:15<01:36,  5.41it/s, 51 steps of size 4.83e-02. acc. prob=0.78]

warmup:  31%|███       | 229/750 [02:15<01:52,  4.62it/s, 127 steps of size 4.95e-02. acc. prob=0.78]

warmup:  31%|███       | 230/750 [02:15<01:42,  5.08it/s, 63 steps of size 5.97e-02. acc. prob=0.78]

warmup:  31%|███       | 232/750 [02:16<02:07,  4.07it/s, 255 steps of size 1.90e-02. acc. prob=0.77]

warmup:  31%|███       | 233/750 [02:17<02:47,  3.08it/s, 255 steps of size 2.63e-02. acc. prob=0.77]

warmup:  31%|███       | 234/750 [02:17<03:21,  2.56it/s, 255 steps of size 3.89e-02. acc. prob=0.78]

warmup:  31%|███▏      | 235/750 [02:17<03:04,  2.79it/s, 127 steps of size 5.58e-02. acc. prob=0.78]

warmup:  31%|███▏      | 236/750 [02:18<02:33,  3.35it/s, 63 steps of size 7.18e-02. acc. prob=0.78]

warmup:  32%|███▏      | 237/750 [02:18<02:09,  3.96it/s, 63 steps of size 9.42e-02. acc. prob=0.78]

warmup:  32%|███▏      | 238/750 [02:18<01:51,  4.57it/s, 63 steps of size 1.16e-01. acc. prob=0.78]

warmup:  32%|███▏      | 241/750 [02:18<01:04,  7.94it/s, 31 steps of size 7.97e-02. acc. prob=0.78]

warmup:  32%|███▏      | 242/750 [02:18<01:05,  7.78it/s, 63 steps of size 5.54e-02. acc. prob=0.78]

warmup:  32%|███▏      | 243/750 [02:18<01:22,  6.14it/s, 127 steps of size 8.04e-02. acc. prob=0.78]

warmup:  33%|███▎      | 245/750 [02:19<01:23,  6.03it/s, 127 steps of size 3.22e-02. acc. prob=0.78]

warmup:  33%|███▎      | 246/750 [02:19<01:35,  5.27it/s, 127 steps of size 9.35e-03. acc. prob=0.77]

warmup:  33%|███▎      | 247/750 [02:20<03:24,  2.46it/s, 511 steps of size 1.36e-02. acc. prob=0.77]

warmup:  33%|███▎      | 248/750 [02:21<03:41,  2.26it/s, 255 steps of size 1.66e-02. acc. prob=0.77]

warmup:  33%|███▎      | 249/750 [02:21<03:55,  2.13it/s, 255 steps of size 2.42e-02. acc. prob=0.78]

warmup:  33%|███▎      | 250/750 [02:21<03:27,  2.41it/s, 127 steps of size 1.04e-02. acc. prob=0.77]

warmup:  33%|███▎      | 251/750 [02:23<05:04,  1.64it/s, 511 steps of size 1.37e-01. acc. prob=0.77]

warmup:  34%|███▎      | 253/750 [02:23<03:49,  2.17it/s, 255 steps of size 2.74e-02. acc. prob=0.77]

warmup:  34%|███▍      | 254/750 [02:23<03:26,  2.40it/s, 127 steps of size 3.66e-02. acc. prob=0.77]

warmup:  34%|███▍      | 255/750 [02:24<03:07,  2.63it/s, 127 steps of size 4.31e-02. acc. prob=0.77]

warmup:  34%|███▍      | 256/750 [02:24<02:30,  3.28it/s, 47 steps of size 2.05e-02. acc. prob=0.77]

warmup:  34%|███▍      | 257/750 [02:24<03:03,  2.68it/s, 255 steps of size 3.26e-02. acc. prob=0.77]

warmup:  34%|███▍      | 258/750 [02:25<02:49,  2.90it/s, 127 steps of size 4.67e-02. acc. prob=0.77]

warmup:  35%|███▍      | 259/750 [02:25<02:39,  3.07it/s, 127 steps of size 2.31e-02. acc. prob=0.77]

warmup:  35%|███▍      | 260/750 [02:25<03:11,  2.56it/s, 255 steps of size 3.05e-02. acc. prob=0.77]

warmup:  35%|███▍      | 261/750 [02:26<02:54,  2.80it/s, 127 steps of size 5.68e-02. acc. prob=0.77]

warmup:  35%|███▌      | 263/750 [02:28<05:47,  1.40it/s, 1023 steps of size 9.99e-03. acc. prob=0.77]

warmup:  35%|███▌      | 264/750 [02:29<06:33,  1.24it/s, 511 steps of size 1.87e-02. acc. prob=0.77]

warmup:  35%|███▌      | 265/750 [02:30<06:00,  1.35it/s, 255 steps of size 3.47e-02. acc. prob=0.77]

warmup:  36%|███▌      | 267/750 [02:30<03:59,  2.02it/s, 127 steps of size 5.31e-02. acc. prob=0.77]

warmup:  36%|███▌      | 269/750 [02:31<04:09,  1.93it/s, 511 steps of size 1.54e-02. acc. prob=0.77]

warmup:  36%|███▌      | 270/750 [02:32<04:12,  1.90it/s, 255 steps of size 2.58e-02. acc. prob=0.77]

warmup:  36%|███▌      | 271/750 [02:32<03:43,  2.14it/s, 127 steps of size 4.27e-02. acc. prob=0.77]

warmup:  36%|███▋      | 272/750 [02:32<03:20,  2.39it/s, 127 steps of size 7.35e-02. acc. prob=0.77]

warmup:  36%|███▋      | 273/750 [02:32<02:45,  2.89it/s, 63 steps of size 1.13e-01. acc. prob=0.78]

warmup:  37%|███▋      | 275/750 [02:32<01:45,  4.52it/s, 19 steps of size 4.65e-02. acc. prob=0.77]

warmup:  37%|███▋      | 276/750 [02:33<01:35,  4.95it/s, 63 steps of size 7.72e-02. acc. prob=0.78]

warmup:  37%|███▋      | 278/750 [02:33<01:16,  6.16it/s, 63 steps of size 1.46e-01. acc. prob=0.78]

warmup:  37%|███▋      | 280/750 [02:33<00:58,  8.03it/s, 21 steps of size 5.82e-02. acc. prob=0.78]

warmup:  37%|███▋      | 281/750 [02:33<01:00,  7.80it/s, 63 steps of size 9.79e-02. acc. prob=0.78]

warmup:  38%|███▊      | 282/750 [02:33<01:02,  7.48it/s, 63 steps of size 1.28e-01. acc. prob=0.78]

warmup:  38%|███▊      | 284/750 [02:33<00:50,  9.18it/s, 31 steps of size 1.50e-01. acc. prob=0.78]

warmup:  38%|███▊      | 286/750 [02:34<01:24,  5.52it/s, 255 steps of size 4.33e-02. acc. prob=0.78]

warmup:  38%|███▊      | 287/750 [02:34<01:32,  5.00it/s, 127 steps of size 7.32e-02. acc. prob=0.78]

warmup:  38%|███▊      | 288/750 [02:34<01:25,  5.38it/s, 63 steps of size 1.20e-01. acc. prob=0.78]

warmup:  39%|███▊      | 290/750 [02:34<01:04,  7.18it/s, 31 steps of size 1.69e-01. acc. prob=0.78]

warmup:  39%|███▉      | 293/750 [02:35<00:50,  8.96it/s, 63 steps of size 1.18e-01. acc. prob=0.78]

warmup:  39%|███▉      | 296/750 [02:35<00:38, 11.95it/s, 11 steps of size 3.82e-02. acc. prob=0.78]

warmup:  40%|███▉      | 297/750 [02:35<00:51,  8.80it/s, 127 steps of size 6.04e-02. acc. prob=0.78]

warmup:  40%|███▉      | 298/750 [02:35<00:53,  8.49it/s, 63 steps of size 8.59e-02. acc. prob=0.78]

warmup:  40%|███▉      | 299/750 [02:35<00:54,  8.22it/s, 63 steps of size 1.38e-01. acc. prob=0.78]

warmup:  40%|████      | 301/750 [02:36<01:02,  7.16it/s, 127 steps of size 6.68e-02. acc. prob=0.78]

warmup:  40%|████      | 302/750 [02:36<01:02,  7.21it/s, 63 steps of size 9.11e-02. acc. prob=0.78]

warmup:  41%|████      | 304/750 [02:36<00:50,  8.91it/s, 31 steps of size 6.21e-02. acc. prob=0.78]

warmup:  41%|████      | 305/750 [02:36<00:52,  8.53it/s, 63 steps of size 9.60e-02. acc. prob=0.78]

warmup:  41%|████      | 306/750 [02:36<00:53,  8.23it/s, 63 steps of size 1.49e-01. acc. prob=0.78]

warmup:  41%|████      | 308/750 [02:36<00:48,  9.15it/s, 63 steps of size 1.29e-01. acc. prob=0.78]

warmup:  41%|████▏     | 310/750 [02:37<00:57,  7.62it/s, 127 steps of size 4.59e-02. acc. prob=0.78]

warmup:  41%|████▏     | 311/750 [02:37<00:58,  7.57it/s, 63 steps of size 6.99e-02. acc. prob=0.78]

warmup:  42%|████▏     | 312/750 [02:37<00:58,  7.53it/s, 63 steps of size 1.06e-01. acc. prob=0.78]

warmup:  42%|████▏     | 315/750 [02:37<00:42, 10.29it/s, 63 steps of size 9.76e-02. acc. prob=0.78]

warmup:  42%|████▏     | 316/750 [02:37<00:45,  9.57it/s, 63 steps of size 8.67e-02. acc. prob=0.78]

warmup:  42%|████▏     | 318/750 [02:38<00:56,  7.59it/s, 127 steps of size 5.31e-02. acc. prob=0.78]

warmup:  43%|████▎     | 319/750 [02:38<00:57,  7.54it/s, 63 steps of size 8.01e-02. acc. prob=0.78]

warmup:  43%|████▎     | 320/750 [02:38<00:57,  7.51it/s, 63 steps of size 4.66e-02. acc. prob=0.78]

warmup:  43%|████▎     | 321/750 [02:38<00:57,  7.47it/s, 63 steps of size 6.44e-02. acc. prob=0.78]

warmup:  43%|████▎     | 322/750 [02:38<00:57,  7.42it/s, 63 steps of size 3.66e-02. acc. prob=0.78]

warmup:  43%|████▎     | 323/750 [02:39<01:13,  5.82it/s, 127 steps of size 5.40e-02. acc. prob=0.78]

warmup:  43%|████▎     | 324/750 [02:39<01:09,  6.16it/s, 63 steps of size 2.88e-02. acc. prob=0.78]

warmup:  43%|████▎     | 325/750 [02:39<01:55,  3.68it/s, 255 steps of size 3.21e-02. acc. prob=0.78]

warmup:  43%|████▎     | 326/750 [02:40<01:55,  3.66it/s, 127 steps of size 2.88e-02. acc. prob=0.78]

warmup:  44%|████▎     | 327/750 [02:40<01:55,  3.65it/s, 127 steps of size 4.27e-02. acc. prob=0.78]

warmup:  44%|████▎     | 328/750 [02:40<01:55,  3.65it/s, 127 steps of size 5.19e-02. acc. prob=0.78]

warmup:  44%|████▍     | 329/750 [02:40<01:55,  3.65it/s, 127 steps of size 4.31e-02. acc. prob=0.78]

warmup:  44%|████▍     | 330/750 [02:40<01:37,  4.29it/s, 63 steps of size 6.07e-02. acc. prob=0.78]

warmup:  44%|████▍     | 331/750 [02:41<01:25,  4.89it/s, 63 steps of size 3.86e-02. acc. prob=0.78]

warmup:  44%|████▍     | 332/750 [02:41<01:34,  4.44it/s, 127 steps of size 5.47e-02. acc. prob=0.78]

warmup:  44%|████▍     | 333/750 [02:41<01:23,  5.01it/s, 63 steps of size 6.09e-02. acc. prob=0.78]

warmup:  45%|████▍     | 334/750 [02:41<01:15,  5.54it/s, 63 steps of size 4.62e-02. acc. prob=0.78]

warmup:  45%|████▍     | 335/750 [02:41<01:26,  4.79it/s, 127 steps of size 5.59e-02. acc. prob=0.78]

warmup:  45%|████▍     | 336/750 [02:42<01:17,  5.34it/s, 63 steps of size 5.73e-02. acc. prob=0.78]

warmup:  45%|████▍     | 337/750 [02:42<01:11,  5.81it/s, 63 steps of size 5.45e-02. acc. prob=0.78]

warmup:  45%|████▌     | 338/750 [02:42<01:06,  6.21it/s, 63 steps of size 4.86e-02. acc. prob=0.78]

warmup:  45%|████▌     | 339/750 [02:42<01:19,  5.14it/s, 127 steps of size 7.08e-02. acc. prob=0.78]

warmup:  45%|████▌     | 340/750 [02:42<01:12,  5.64it/s, 63 steps of size 6.11e-02. acc. prob=0.78]

warmup:  45%|████▌     | 341/750 [02:42<01:07,  6.05it/s, 63 steps of size 7.86e-02. acc. prob=0.78]

warmup:  46%|████▌     | 342/750 [02:43<01:04,  6.35it/s, 63 steps of size 8.36e-02. acc. prob=0.78]

warmup:  46%|████▌     | 343/750 [02:43<01:01,  6.61it/s, 63 steps of size 1.91e-02. acc. prob=0.78]

warmup:  46%|████▌     | 344/750 [02:43<01:49,  3.70it/s, 255 steps of size 2.74e-02. acc. prob=0.78]

warmup:  46%|████▌     | 345/750 [02:44<01:49,  3.69it/s, 127 steps of size 2.39e-02. acc. prob=0.78]

warmup:  46%|████▌     | 346/750 [02:44<01:49,  3.68it/s, 127 steps of size 3.47e-02. acc. prob=0.78]

warmup:  46%|████▋     | 347/750 [02:44<01:31,  4.42it/s, 55 steps of size 2.38e-02. acc. prob=0.78]

warmup:  46%|████▋     | 348/750 [02:44<01:36,  4.16it/s, 127 steps of size 3.29e-02. acc. prob=0.78]

warmup:  47%|████▋     | 349/750 [02:44<01:40,  3.97it/s, 127 steps of size 4.39e-02. acc. prob=0.78]

warmup:  47%|████▋     | 351/750 [02:46<02:37,  2.54it/s, 511 steps of size 1.49e-02. acc. prob=0.78]

warmup:  47%|████▋     | 352/750 [02:46<02:52,  2.31it/s, 255 steps of size 2.08e-02. acc. prob=0.78]

warmup:  47%|████▋     | 353/750 [02:47<03:04,  2.15it/s, 255 steps of size 2.99e-02. acc. prob=0.78]

warmup:  47%|████▋     | 354/750 [02:47<02:28,  2.67it/s, 63 steps of size 1.69e-02. acc. prob=0.78]

warmup:  47%|████▋     | 355/750 [02:47<02:42,  2.43it/s, 255 steps of size 2.43e-02. acc. prob=0.78]

warmup:  47%|████▋     | 356/750 [02:48<02:24,  2.72it/s, 127 steps of size 3.46e-02. acc. prob=0.78]

warmup:  48%|████▊     | 357/750 [02:48<02:11,  3.00it/s, 127 steps of size 3.07e-02. acc. prob=0.78]

warmup:  48%|████▊     | 358/750 [02:48<02:01,  3.23it/s, 127 steps of size 3.39e-02. acc. prob=0.78]

warmup:  48%|████▊     | 359/750 [02:48<01:52,  3.47it/s, 127 steps of size 3.46e-02. acc. prob=0.78]

warmup:  48%|████▊     | 360/750 [02:49<01:46,  3.66it/s, 127 steps of size 4.83e-02. acc. prob=0.78]

warmup:  48%|████▊     | 361/750 [02:49<01:41,  3.85it/s, 127 steps of size 2.33e-02. acc. prob=0.78]

warmup:  48%|████▊     | 362/750 [02:49<02:04,  3.12it/s, 255 steps of size 2.73e-02. acc. prob=0.78]

warmup:  48%|████▊     | 363/750 [02:49<01:54,  3.39it/s, 127 steps of size 3.60e-02. acc. prob=0.78]

warmup:  49%|████▊     | 364/750 [02:50<01:46,  3.63it/s, 127 steps of size 4.15e-02. acc. prob=0.78]

warmup:  49%|████▊     | 365/750 [02:50<01:41,  3.80it/s, 127 steps of size 5.74e-02. acc. prob=0.78]

warmup:  49%|████▉     | 367/750 [02:50<01:37,  3.93it/s, 255 steps of size 2.08e-02. acc. prob=0.78]

warmup:  49%|████▉     | 368/750 [02:51<01:56,  3.27it/s, 255 steps of size 2.90e-02. acc. prob=0.78]

warmup:  49%|████▉     | 369/750 [02:51<01:49,  3.49it/s, 127 steps of size 4.08e-02. acc. prob=0.78]

warmup:  49%|████▉     | 370/750 [02:51<01:43,  3.69it/s, 127 steps of size 5.29e-02. acc. prob=0.78]

warmup:  49%|████▉     | 371/750 [02:52<01:50,  3.43it/s, 191 steps of size 6.43e-02. acc. prob=0.78]

warmup:  50%|████▉     | 372/750 [02:52<01:31,  4.15it/s, 63 steps of size 8.53e-02. acc. prob=0.78]

warmup:  50%|████▉     | 373/750 [02:52<01:17,  4.89it/s, 63 steps of size 7.32e-02. acc. prob=0.78]

warmup:  50%|████▉     | 374/750 [02:52<01:06,  5.62it/s, 63 steps of size 9.41e-02. acc. prob=0.78]

warmup:  50%|█████     | 375/750 [02:52<00:59,  6.28it/s, 63 steps of size 3.64e-02. acc. prob=0.78]

warmup:  50%|█████     | 376/750 [02:52<01:07,  5.55it/s, 127 steps of size 4.94e-02. acc. prob=0.78]

warmup:  50%|█████     | 377/750 [02:53<01:38,  3.79it/s, 255 steps of size 6.73e-02. acc. prob=0.78]

warmup:  50%|█████     | 378/750 [02:53<01:21,  4.56it/s, 63 steps of size 7.68e-02. acc. prob=0.78]

warmup:  51%|█████     | 379/750 [02:53<01:09,  5.32it/s, 63 steps of size 2.69e-02. acc. prob=0.78]

warmup:  51%|█████     | 380/750 [02:53<01:14,  4.99it/s, 127 steps of size 3.73e-02. acc. prob=0.78]

warmup:  51%|█████     | 381/750 [02:54<01:17,  4.79it/s, 127 steps of size 4.85e-02. acc. prob=0.78]

warmup:  51%|█████     | 382/750 [02:54<01:19,  4.66it/s, 127 steps of size 6.70e-02. acc. prob=0.78]

warmup:  51%|█████     | 383/750 [02:54<01:07,  5.42it/s, 63 steps of size 8.49e-02. acc. prob=0.78]

warmup:  51%|█████     | 384/750 [02:54<00:59,  6.12it/s, 63 steps of size 9.60e-02. acc. prob=0.78]

warmup:  51%|█████▏    | 385/750 [02:54<00:54,  6.73it/s, 63 steps of size 2.68e-02. acc. prob=0.78]

warmup:  51%|█████▏    | 386/750 [02:55<01:28,  4.13it/s, 255 steps of size 3.66e-02. acc. prob=0.78]

warmup:  52%|█████▏    | 387/750 [02:55<01:26,  4.19it/s, 127 steps of size 4.67e-02. acc. prob=0.78]

warmup:  52%|█████▏    | 388/750 [02:55<01:25,  4.25it/s, 127 steps of size 6.31e-02. acc. prob=0.78]

warmup:  52%|█████▏    | 389/750 [02:55<01:11,  5.03it/s, 63 steps of size 7.46e-02. acc. prob=0.78]

warmup:  52%|█████▏    | 390/750 [02:55<01:02,  5.76it/s, 63 steps of size 1.01e-01. acc. prob=0.78]

warmup:  52%|█████▏    | 391/750 [02:55<00:55,  6.42it/s, 63 steps of size 8.35e-02. acc. prob=0.78]

warmup:  52%|█████▏    | 392/750 [02:55<00:51,  6.99it/s, 63 steps of size 1.10e-01. acc. prob=0.78]

warmup:  53%|█████▎    | 394/750 [02:56<00:36,  9.72it/s, 31 steps of size 1.68e-01. acc. prob=0.78]

warmup:  53%|█████▎    | 396/750 [02:56<00:30, 11.78it/s, 31 steps of size 9.52e-02. acc. prob=0.78]

warmup:  53%|█████▎    | 398/750 [02:56<00:26, 13.42it/s, 31 steps of size 1.72e-01. acc. prob=0.78]

warmup:  53%|█████▎    | 400/750 [02:56<00:23, 14.60it/s, 31 steps of size 1.69e-01. acc. prob=0.78]

warmup:  54%|█████▎    | 403/750 [02:56<00:24, 14.45it/s, 63 steps of size 1.23e-01. acc. prob=0.78]

warmup:  54%|█████▍    | 405/750 [02:56<00:22, 15.27it/s, 31 steps of size 2.19e-01. acc. prob=0.78]

warmup:  54%|█████▍    | 407/750 [02:56<00:22, 14.96it/s, 63 steps of size 1.07e-01. acc. prob=0.78]

warmup:  55%|█████▍    | 409/750 [02:57<00:24, 13.88it/s, 63 steps of size 1.13e-01. acc. prob=0.78]

warmup:  55%|█████▍    | 412/750 [02:57<00:28, 11.99it/s, 127 steps of size 5.45e-02. acc. prob=0.78]

warmup:  55%|█████▌    | 413/750 [02:57<00:29, 11.35it/s, 63 steps of size 7.20e-02. acc. prob=0.78]

warmup:  55%|█████▌    | 414/750 [02:57<00:31, 10.79it/s, 63 steps of size 8.95e-02. acc. prob=0.78]

warmup:  55%|█████▌    | 415/750 [02:57<00:32, 10.31it/s, 63 steps of size 4.24e-02. acc. prob=0.78]

warmup:  55%|█████▌    | 416/750 [02:57<00:42,  7.89it/s, 127 steps of size 4.29e-02. acc. prob=0.78]

warmup:  56%|█████▌    | 417/750 [02:58<00:50,  6.60it/s, 127 steps of size 5.70e-02. acc. prob=0.78]

warmup:  56%|█████▌    | 418/750 [02:58<00:47,  7.02it/s, 63 steps of size 4.97e-02. acc. prob=0.78]

warmup:  56%|█████▌    | 419/750 [02:58<00:54,  6.04it/s, 127 steps of size 5.78e-02. acc. prob=0.78]

warmup:  56%|█████▌    | 420/750 [02:58<00:49,  6.62it/s, 63 steps of size 5.51e-02. acc. prob=0.78]

warmup:  56%|█████▌    | 421/750 [02:58<00:46,  7.12it/s, 63 steps of size 7.16e-02. acc. prob=0.78]

warmup:  56%|█████▋    | 422/750 [02:58<00:43,  7.53it/s, 63 steps of size 3.09e-02. acc. prob=0.78]

warmup:  56%|█████▋    | 423/750 [02:59<00:52,  6.22it/s, 127 steps of size 4.06e-02. acc. prob=0.78]

warmup:  57%|█████▋    | 424/750 [02:59<00:58,  5.53it/s, 127 steps of size 5.42e-02. acc. prob=0.78]

warmup:  57%|█████▋    | 425/750 [02:59<00:52,  6.22it/s, 63 steps of size 7.10e-02. acc. prob=0.78]

warmup:  57%|█████▋    | 426/750 [02:59<00:47,  6.81it/s, 63 steps of size 8.52e-02. acc. prob=0.78]

warmup:  57%|█████▋    | 427/750 [02:59<00:44,  7.31it/s, 63 steps of size 1.13e-01. acc. prob=0.78]

warmup:  57%|█████▋    | 429/750 [02:59<00:41,  7.68it/s, 127 steps of size 4.81e-02. acc. prob=0.78]

warmup:  57%|█████▋    | 430/750 [03:00<00:49,  6.45it/s, 127 steps of size 6.34e-02. acc. prob=0.78]

warmup:  58%|█████▊    | 432/750 [03:00<01:00,  5.26it/s, 255 steps of size 3.04e-02. acc. prob=0.78]

warmup:  58%|█████▊    | 433/750 [03:00<01:03,  5.01it/s, 127 steps of size 3.97e-02. acc. prob=0.78]

warmup:  58%|█████▊    | 434/750 [03:01<01:05,  4.84it/s, 127 steps of size 1.48e-02. acc. prob=0.78]

warmup:  58%|█████▊    | 435/750 [03:01<01:25,  3.67it/s, 255 steps of size 1.92e-02. acc. prob=0.78]

warmup:  58%|█████▊    | 436/750 [03:01<01:41,  3.09it/s, 255 steps of size 2.53e-02. acc. prob=0.78]

warmup:  58%|█████▊    | 437/750 [03:02<01:33,  3.36it/s, 127 steps of size 3.26e-02. acc. prob=0.78]

warmup:  58%|█████▊    | 438/750 [03:02<01:26,  3.59it/s, 127 steps of size 4.32e-02. acc. prob=0.78]

warmup:  59%|█████▊    | 439/750 [03:02<01:22,  3.77it/s, 127 steps of size 4.85e-02. acc. prob=0.78]

warmup:  59%|█████▊    | 440/750 [03:02<01:19,  3.91it/s, 127 steps of size 6.22e-02. acc. prob=0.78]

warmup:  59%|█████▉    | 441/750 [03:03<01:06,  4.66it/s, 63 steps of size 7.93e-02. acc. prob=0.78]

warmup:  59%|█████▉    | 443/750 [03:03<00:43,  7.00it/s, 31 steps of size 1.16e-01. acc. prob=0.79]

warmup:  59%|█████▉    | 445/750 [03:03<00:36,  8.28it/s, 63 steps of size 8.51e-02. acc. prob=0.78]

warmup:  59%|█████▉    | 446/750 [03:03<00:36,  8.35it/s, 63 steps of size 1.03e-01. acc. prob=0.78]

warmup:  60%|█████▉    | 448/750 [03:03<00:28, 10.49it/s, 31 steps of size 1.72e-01. acc. prob=0.79]

warmup:  60%|██████    | 450/750 [03:03<00:24, 12.27it/s, 31 steps of size 1.31e-01. acc. prob=0.79]

warmup:  61%|██████    | 454/750 [03:03<00:19, 15.48it/s, 63 steps of size 8.75e-02. acc. prob=0.78]

warmup:  61%|██████    | 456/750 [03:03<00:18, 16.03it/s, 31 steps of size 1.98e-01. acc. prob=0.78]

warmup:  61%|██████    | 458/750 [03:04<00:24, 11.85it/s, 127 steps of size 7.58e-02. acc. prob=0.78]

warmup:  61%|██████    | 459/750 [03:04<00:25, 11.22it/s, 63 steps of size 1.08e-01. acc. prob=0.78]

warmup:  61%|██████▏   | 461/750 [03:04<00:22, 12.72it/s, 31 steps of size 3.59e-01. acc. prob=0.78]

warmup:  62%|██████▏   | 463/750 [03:04<00:25, 11.05it/s, 127 steps of size 5.41e-02. acc. prob=0.78]

warmup:  62%|██████▏   | 464/750 [03:04<00:32,  8.68it/s, 127 steps of size 1.02e-01. acc. prob=0.78]

warmup:  62%|██████▏   | 468/750 [03:05<00:27, 10.30it/s, 127 steps of size 5.76e-02. acc. prob=0.78]

warmup:  63%|██████▎   | 469/750 [03:05<00:27, 10.05it/s, 63 steps of size 1.08e-01. acc. prob=0.78]

warmup:  63%|██████▎   | 471/750 [03:05<00:23, 11.63it/s, 31 steps of size 2.49e-01. acc. prob=0.78]

warmup:  63%|██████▎   | 473/750 [03:05<00:22, 12.27it/s, 63 steps of size 1.03e-01. acc. prob=0.78]

warmup:  63%|██████▎   | 475/750 [03:05<00:20, 13.58it/s, 31 steps of size 4.78e-02. acc. prob=0.78]

warmup:  63%|██████▎   | 476/750 [03:05<00:22, 12.42it/s, 63 steps of size 8.06e-02. acc. prob=0.78]

warmup:  64%|██████▎   | 477/750 [03:05<00:23, 11.50it/s, 63 steps of size 1.43e-01. acc. prob=0.78]

warmup:  64%|██████▍   | 479/750 [03:06<00:22, 11.83it/s, 63 steps of size 9.33e-02. acc. prob=0.78]

warmup:  64%|██████▍   | 480/750 [03:06<00:24, 11.07it/s, 63 steps of size 1.40e-01. acc. prob=0.78]

warmup:  64%|██████▍   | 482/750 [03:06<00:29,  9.18it/s, 127 steps of size 6.29e-02. acc. prob=0.78]

warmup:  64%|██████▍   | 483/750 [03:06<00:35,  7.49it/s, 127 steps of size 1.08e-01. acc. prob=0.78]

warmup:  65%|██████▍   | 484/750 [03:06<00:34,  7.78it/s, 63 steps of size 1.23e-01. acc. prob=0.78]

warmup:  65%|██████▍   | 486/750 [03:07<00:29,  9.03it/s, 63 steps of size 5.00e-02. acc. prob=0.78]

warmup:  65%|██████▍   | 487/750 [03:07<00:29,  9.00it/s, 63 steps of size 8.53e-02. acc. prob=0.78]

warmup:  65%|██████▌   | 488/750 [03:07<00:29,  8.97it/s, 63 steps of size 1.41e-01. acc. prob=0.78]

warmup:  65%|██████▌   | 490/750 [03:07<00:23, 10.88it/s, 31 steps of size 3.37e-02. acc. prob=0.78]

warmup:  65%|██████▌   | 491/750 [03:07<00:32,  8.07it/s, 127 steps of size 5.43e-02. acc. prob=0.78]

warmup:  66%|██████▌   | 492/750 [03:07<00:31,  8.26it/s, 63 steps of size 9.01e-02. acc. prob=0.78]

warmup:  66%|██████▌   | 494/750 [03:07<00:24, 10.62it/s, 31 steps of size 1.54e-01. acc. prob=0.79]

warmup:  66%|██████▋   | 497/750 [03:07<00:18, 13.97it/s, 31 steps of size 1.99e-01. acc. prob=0.79]

warmup:  67%|██████▋   | 500/750 [03:08<00:15, 15.88it/s, 63 steps of size 1.02e-01. acc. prob=0.78]

sample:  67%|██████▋   | 502/750 [03:08<00:15, 16.42it/s, 31 steps of size 1.02e-01. acc. prob=0.95]

sample:  67%|██████▋   | 504/750 [03:08<00:14, 16.86it/s, 31 steps of size 1.02e-01. acc. prob=0.97]

sample:  67%|██████▋   | 506/750 [03:08<00:14, 17.18it/s, 31 steps of size 1.02e-01. acc. prob=0.87]

sample:  68%|██████▊   | 508/750 [03:08<00:13, 17.43it/s, 31 steps of size 1.02e-01. acc. prob=0.85]

sample:  68%|██████▊   | 511/750 [03:08<00:12, 18.72it/s, 34 steps of size 1.02e-01. acc. prob=0.78]

sample:  68%|██████▊   | 513/750 [03:08<00:16, 14.46it/s, 95 steps of size 1.02e-01. acc. prob=0.79]

sample:  69%|██████▊   | 515/750 [03:09<00:15, 15.28it/s, 31 steps of size 1.02e-01. acc. prob=0.75]

sample:  69%|██████▉   | 517/750 [03:09<00:14, 15.96it/s, 31 steps of size 1.02e-01. acc. prob=0.76]

sample:  69%|██████▉   | 518/750 [03:09<00:16, 14.07it/s, 63 steps of size 1.02e-01. acc. prob=0.71]

sample:  69%|██████▉   | 519/750 [03:09<00:23,  9.87it/s, 127 steps of size 1.02e-01. acc. prob=0.70]

sample:  69%|██████▉   | 521/750 [03:09<00:19, 11.77it/s, 31 steps of size 1.02e-01. acc. prob=0.66]

sample:  70%|██████▉   | 523/750 [03:09<00:17, 13.34it/s, 31 steps of size 1.02e-01. acc. prob=0.64]

sample:  70%|███████   | 525/750 [03:09<00:17, 12.85it/s, 63 steps of size 1.02e-01. acc. prob=0.67]

sample:  70%|███████   | 526/750 [03:09<00:18, 11.85it/s, 63 steps of size 1.02e-01. acc. prob=0.68]

sample:  70%|███████   | 528/750 [03:10<00:18, 11.88it/s, 63 steps of size 1.02e-01. acc. prob=0.70]

sample:  71%|███████   | 530/750 [03:10<00:16, 13.40it/s, 31 steps of size 1.02e-01. acc. prob=0.72]

sample:  71%|███████   | 532/750 [03:10<00:17, 12.79it/s, 63 steps of size 1.02e-01. acc. prob=0.73]

sample:  71%|███████   | 534/750 [03:10<00:15, 14.06it/s, 31 steps of size 1.02e-01. acc. prob=0.74]

sample:  71%|███████▏  | 536/750 [03:10<00:14, 15.09it/s, 31 steps of size 1.02e-01. acc. prob=0.75]

sample:  72%|███████▏  | 538/750 [03:10<00:13, 15.89it/s, 31 steps of size 1.02e-01. acc. prob=0.76]

sample:  72%|███████▏  | 540/750 [03:10<00:12, 16.49it/s, 31 steps of size 1.02e-01. acc. prob=0.77]

sample:  72%|███████▏  | 541/750 [03:10<00:14, 14.31it/s, 63 steps of size 1.02e-01. acc. prob=0.77]

sample:  72%|███████▏  | 542/750 [03:11<00:16, 12.74it/s, 63 steps of size 1.02e-01. acc. prob=0.78]

sample:  73%|███████▎  | 544/750 [03:11<00:14, 14.17it/s, 31 steps of size 1.02e-01. acc. prob=0.79]

sample:  73%|███████▎  | 546/750 [03:11<00:13, 15.30it/s, 31 steps of size 1.02e-01. acc. prob=0.80]

sample:  73%|███████▎  | 548/750 [03:11<00:14, 14.00it/s, 63 steps of size 1.02e-01. acc. prob=0.80]

sample:  73%|███████▎  | 549/750 [03:11<00:15, 12.62it/s, 63 steps of size 1.02e-01. acc. prob=0.81]

sample:  73%|███████▎  | 550/750 [03:11<00:22,  9.04it/s, 127 steps of size 1.02e-01. acc. prob=0.81]

sample:  73%|███████▎  | 551/750 [03:12<00:27,  7.26it/s, 127 steps of size 1.02e-01. acc. prob=0.81]

sample:  74%|███████▎  | 553/750 [03:12<00:20,  9.55it/s, 31 steps of size 1.02e-01. acc. prob=0.82]

sample:  74%|███████▍  | 555/750 [03:12<00:16, 11.52it/s, 31 steps of size 1.02e-01. acc. prob=0.82]

sample:  74%|███████▍  | 557/750 [03:12<00:14, 13.15it/s, 31 steps of size 1.02e-01. acc. prob=0.83]

sample:  75%|███████▍  | 559/750 [03:12<00:13, 14.44it/s, 31 steps of size 1.02e-01. acc. prob=0.83]

sample:  75%|███████▍  | 561/750 [03:12<00:12, 15.41it/s, 31 steps of size 1.02e-01. acc. prob=0.84]

sample:  75%|███████▌  | 563/750 [03:12<00:11, 16.15it/s, 31 steps of size 1.02e-01. acc. prob=0.84]

sample:  75%|███████▌  | 565/750 [03:13<00:16, 11.56it/s, 127 steps of size 1.02e-01. acc. prob=0.84]

sample:  76%|███████▌  | 567/750 [03:13<00:15, 11.67it/s, 63 steps of size 1.02e-01. acc. prob=0.85]

sample:  76%|███████▌  | 569/750 [03:13<00:13, 13.07it/s, 31 steps of size 1.02e-01. acc. prob=0.85]

sample:  76%|███████▌  | 570/750 [03:13<00:14, 12.06it/s, 63 steps of size 1.02e-01. acc. prob=0.85]

sample:  76%|███████▋  | 572/750 [03:13<00:13, 13.52it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  76%|███████▋  | 573/750 [03:13<00:14, 12.23it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  77%|███████▋  | 575/750 [03:13<00:14, 12.12it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  77%|███████▋  | 577/750 [03:13<00:14, 12.06it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  77%|███████▋  | 578/750 [03:14<00:15, 11.27it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  77%|███████▋  | 579/750 [03:14<00:16, 10.66it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  77%|███████▋  | 581/750 [03:14<00:13, 12.63it/s, 31 steps of size 1.02e-01. acc. prob=0.87]

sample:  78%|███████▊  | 583/750 [03:14<00:11, 14.13it/s, 31 steps of size 1.02e-01. acc. prob=0.87]

sample:  78%|███████▊  | 585/750 [03:14<00:10, 15.24it/s, 31 steps of size 1.02e-01. acc. prob=0.87]

sample:  78%|███████▊  | 586/750 [03:14<00:12, 13.40it/s, 63 steps of size 1.02e-01. acc. prob=0.87]

sample:  78%|███████▊  | 588/750 [03:14<00:10, 14.73it/s, 31 steps of size 1.02e-01. acc. prob=0.87]

sample:  79%|███████▊  | 589/750 [03:14<00:12, 12.92it/s, 63 steps of size 1.02e-01. acc. prob=0.88]

sample:  79%|███████▉  | 591/750 [03:15<00:12, 12.55it/s, 63 steps of size 1.02e-01. acc. prob=0.88]

sample:  79%|███████▉  | 593/750 [03:15<00:11, 14.01it/s, 31 steps of size 1.02e-01. acc. prob=0.88]

sample:  79%|███████▉  | 594/750 [03:15<00:12, 12.57it/s, 63 steps of size 1.02e-01. acc. prob=0.88]

sample:  79%|███████▉  | 596/750 [03:15<00:10, 14.11it/s, 31 steps of size 1.02e-01. acc. prob=0.88]

sample:  80%|███████▉  | 597/750 [03:15<00:12, 12.58it/s, 63 steps of size 1.02e-01. acc. prob=0.88]

sample:  80%|███████▉  | 599/750 [03:15<00:12, 12.33it/s, 63 steps of size 1.02e-01. acc. prob=0.88]

sample:  80%|████████  | 600/750 [03:15<00:13, 11.41it/s, 63 steps of size 1.02e-01. acc. prob=0.88]

sample:  80%|████████  | 602/750 [03:15<00:12, 11.60it/s, 63 steps of size 1.02e-01. acc. prob=0.89]

sample:  80%|████████  | 603/750 [03:16<00:13, 10.90it/s, 63 steps of size 1.02e-01. acc. prob=0.89]

sample:  81%|████████  | 604/750 [03:16<00:14, 10.37it/s, 63 steps of size 1.02e-01. acc. prob=0.89]

sample:  81%|████████  | 606/750 [03:16<00:11, 12.45it/s, 31 steps of size 1.02e-01. acc. prob=0.89]

sample:  81%|████████  | 608/750 [03:16<00:10, 14.03it/s, 31 steps of size 1.02e-01. acc. prob=0.89]

sample:  81%|████████▏ | 610/750 [03:16<00:09, 15.19it/s, 31 steps of size 1.02e-01. acc. prob=0.89]

sample:  82%|████████▏ | 612/750 [03:16<00:08, 16.04it/s, 31 steps of size 1.02e-01. acc. prob=0.89]

sample:  82%|████████▏ | 614/750 [03:16<00:09, 14.38it/s, 63 steps of size 1.02e-01. acc. prob=0.89]

sample:  82%|████████▏ | 616/750 [03:16<00:08, 15.35it/s, 31 steps of size 1.02e-01. acc. prob=0.89]

sample:  82%|████████▏ | 618/750 [03:16<00:08, 16.08it/s, 31 steps of size 1.02e-01. acc. prob=0.90]

sample:  83%|████████▎ | 619/750 [03:17<00:09, 14.03it/s, 63 steps of size 1.02e-01. acc. prob=0.90]

sample:  83%|████████▎ | 621/750 [03:17<00:08, 15.19it/s, 31 steps of size 1.02e-01. acc. prob=0.90]

sample:  83%|████████▎ | 623/750 [03:17<00:09, 13.89it/s, 63 steps of size 1.02e-01. acc. prob=0.90]

sample:  83%|████████▎ | 625/750 [03:17<00:09, 12.78it/s, 63 steps of size 1.02e-01. acc. prob=0.90]

sample:  84%|████████▎ | 627/750 [03:17<00:09, 12.51it/s, 63 steps of size 1.02e-01. acc. prob=0.90]

sample:  84%|████████▍ | 629/750 [03:17<00:09, 12.33it/s, 63 steps of size 1.02e-01. acc. prob=0.90]

sample:  84%|████████▍ | 630/750 [03:17<00:10, 11.53it/s, 63 steps of size 1.02e-01. acc. prob=0.90]

sample:  84%|████████▍ | 631/750 [03:18<00:10, 10.87it/s, 63 steps of size 1.02e-01. acc. prob=0.90]

sample:  84%|████████▍ | 633/750 [03:18<00:11,  9.99it/s, 95 steps of size 1.02e-01. acc. prob=0.90]

sample:  85%|████████▍ | 634/750 [03:18<00:11,  9.74it/s, 63 steps of size 1.02e-01. acc. prob=0.90]

sample:  85%|████████▍ | 635/750 [03:18<00:12,  9.53it/s, 63 steps of size 1.02e-01. acc. prob=0.90]

sample:  85%|████████▍ | 637/750 [03:18<00:09, 11.72it/s, 31 steps of size 1.02e-01. acc. prob=0.90]

sample:  85%|████████▌ | 639/750 [03:18<00:09, 11.32it/s, 63 steps of size 1.02e-01. acc. prob=0.89]

sample:  85%|████████▌ | 641/750 [03:18<00:08, 12.97it/s, 31 steps of size 1.02e-01. acc. prob=0.89]

sample:  86%|████████▌ | 643/750 [03:19<00:07, 14.28it/s, 31 steps of size 1.02e-01. acc. prob=0.89]

sample:  86%|████████▌ | 645/750 [03:19<00:07, 13.44it/s, 63 steps of size 1.02e-01. acc. prob=0.89]

sample:  86%|████████▋ | 647/750 [03:19<00:07, 14.61it/s, 31 steps of size 1.02e-01. acc. prob=0.89]

sample:  86%|████████▋ | 648/750 [03:19<00:08, 11.52it/s, 95 steps of size 1.02e-01. acc. prob=0.89]

sample:  87%|████████▋ | 652/750 [03:19<00:05, 17.27it/s, 15 steps of size 1.02e-01. acc. prob=0.87]

sample:  87%|████████▋ | 654/750 [03:19<00:05, 16.19it/s, 63 steps of size 1.02e-01. acc. prob=0.87]

sample:  87%|████████▋ | 655/750 [03:19<00:06, 14.21it/s, 63 steps of size 1.02e-01. acc. prob=0.87]

sample:  88%|████████▊ | 657/750 [03:20<00:07, 11.93it/s, 95 steps of size 1.02e-01. acc. prob=0.86]

sample:  88%|████████▊ | 658/750 [03:20<00:08, 11.21it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  88%|████████▊ | 659/750 [03:20<00:08, 10.49it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  88%|████████▊ | 660/750 [03:20<00:09,  9.90it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  88%|████████▊ | 661/750 [03:20<00:09,  9.57it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  88%|████████▊ | 662/750 [03:20<00:09,  9.22it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  88%|████████▊ | 663/750 [03:20<00:09,  8.82it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  89%|████████▊ | 664/750 [03:20<00:10,  8.57it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  89%|████████▊ | 665/750 [03:21<00:10,  8.36it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  89%|████████▉ | 666/750 [03:21<00:09,  8.42it/s, 63 steps of size 1.02e-01. acc. prob=0.85]

sample:  89%|████████▉ | 667/750 [03:21<00:09,  8.49it/s, 63 steps of size 1.02e-01. acc. prob=0.85]

sample:  89%|████████▉ | 669/750 [03:21<00:07, 10.99it/s, 34 steps of size 1.02e-01. acc. prob=0.85]

sample:  89%|████████▉ | 670/750 [03:21<00:07, 10.31it/s, 63 steps of size 1.02e-01. acc. prob=0.85]

sample:  89%|████████▉ | 671/750 [03:21<00:07,  9.88it/s, 63 steps of size 1.02e-01. acc. prob=0.85]

sample:  90%|████████▉ | 673/750 [03:21<00:08,  9.38it/s, 95 steps of size 1.02e-01. acc. prob=0.85]

sample:  90%|█████████ | 675/750 [03:22<00:07, 10.17it/s, 63 steps of size 1.02e-01. acc. prob=0.85]

sample:  90%|█████████ | 677/750 [03:22<00:06, 10.69it/s, 63 steps of size 1.02e-01. acc. prob=0.85]

sample:  91%|█████████ | 679/750 [03:22<00:05, 12.35it/s, 31 steps of size 1.02e-01. acc. prob=0.85]

sample:  91%|█████████ | 681/750 [03:22<00:05, 13.73it/s, 31 steps of size 1.02e-01. acc. prob=0.85]

sample:  91%|█████████ | 682/750 [03:22<00:05, 12.43it/s, 63 steps of size 1.02e-01. acc. prob=0.85]

sample:  91%|█████████ | 683/750 [03:22<00:07,  8.91it/s, 127 steps of size 1.02e-01. acc. prob=0.85]

sample:  91%|█████████▏| 685/750 [03:22<00:05, 10.99it/s, 31 steps of size 1.02e-01. acc. prob=0.85]

sample:  91%|█████████▏| 686/750 [03:23<00:06, 10.45it/s, 63 steps of size 1.02e-01. acc. prob=0.85]

sample:  92%|█████████▏| 688/750 [03:23<00:05, 12.34it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  92%|█████████▏| 690/750 [03:23<00:04, 13.75it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  92%|█████████▏| 692/750 [03:23<00:03, 14.90it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  92%|█████████▏| 693/750 [03:23<00:04, 11.52it/s, 95 steps of size 1.02e-01. acc. prob=0.86]

sample:  93%|█████████▎| 694/750 [03:23<00:05, 10.82it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  93%|█████████▎| 696/750 [03:23<00:05,  9.97it/s, 95 steps of size 1.02e-01. acc. prob=0.86]

sample:  93%|█████████▎| 698/750 [03:24<00:04, 11.82it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  93%|█████████▎| 700/750 [03:24<00:03, 13.34it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  94%|█████████▎| 702/750 [03:24<00:03, 14.54it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  94%|█████████▍| 704/750 [03:24<00:02, 15.47it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  94%|█████████▍| 706/750 [03:24<00:02, 16.17it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  94%|█████████▍| 708/750 [03:24<00:02, 16.58it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  95%|█████████▍| 709/750 [03:24<00:03, 12.47it/s, 95 steps of size 1.02e-01. acc. prob=0.86]

sample:  95%|█████████▍| 710/750 [03:24<00:03, 10.11it/s, 95 steps of size 1.02e-01. acc. prob=0.86]

sample:  95%|█████████▍| 711/750 [03:25<00:03,  9.79it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  95%|█████████▌| 713/750 [03:25<00:03, 11.93it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  95%|█████████▌| 715/750 [03:25<00:02, 11.93it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  96%|█████████▌| 717/750 [03:25<00:02, 13.44it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  96%|█████████▌| 719/750 [03:25<00:02, 14.66it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  96%|█████████▌| 721/750 [03:25<00:01, 15.59it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  96%|█████████▋| 723/750 [03:25<00:01, 16.29it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  97%|█████████▋| 725/750 [03:25<00:01, 16.79it/s, 31 steps of size 1.02e-01. acc. prob=0.86]

sample:  97%|█████████▋| 728/750 [03:26<00:01, 15.78it/s, 63 steps of size 1.02e-01. acc. prob=0.86]

sample:  98%|█████████▊| 736/750 [03:26<00:00, 30.38it/s, 7 steps of size 1.02e-01. acc. prob=0.83]

sample:  99%|█████████▉| 745/750 [03:26<00:00, 43.35it/s, 31 steps of size 1.02e-01. acc. prob=0.80]

sample: 100%|█████████▉| 748/750 [03:26<00:00, 35.33it/s, 31 steps of size 1.02e-01. acc. prob=0.80]

sample: 100%|██████████| 750/750 [03:26<00:00, 30.66it/s, 31 steps of size 1.02e-01. acc. prob=0.80]

sample: 100%|██████████| 750/750 [03:26<00:00,  3.63it/s, 31 steps of size 1.02e-01. acc. prob=0.80]

There were 289 divergences after tuning. Increase `target_accept` or reparameterize.

The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details

The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

Let's make sure MCMC converged and can successfully recapitulate the raw data

At an individual level, our model can capture the real data relatively well. How do the summary statistics of the posteriors compare to when we fit using VI?

In [19]:

Copied!





summary_mcmc = az.summary(
    mSimCenteredSampled, var_names=[r".*_mu", r".*_sigma"], filter_vars="regex"
).sort_index()
summary_mcmc
summary_mcmc = az.summary(
    mSimCenteredSampled, var_names=[r".*_mu", r".*_sigma"], filter_vars="regex"
).sort_index()
summary_mcmc

Out[19]:

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
a_1\|participant_id_mu	0.965	0.045	0.89	1	177	241	1.11	0.0033	0.0026
a_1\|participant_id_sigma	0.23	0.04	0.17	0.3	18	30	1.14	0.0097	0.0065
t_1\|participant_id_mu	0.404	0.005	0.4	0.41	33	132	1.10	0.001	0.00062
t_1\|participant_id_sigma	0.005	0.003	0.0015	0.011	12	16	1.24	0.00085	0.00086
v_1\|participant_id_mu	0.27	0.07	0.17	0.38	41	117	1.08	0.0099	0.0083
v_1\|participant_id_sigma	0.31	0.047	0.24	0.39	413	347	1.15	0.0023	0.0021
z_1\|participant_id_mu	0.506	0.008	0.5	0.52	39	226	1.08	0.0012	0.00089
z_1\|participant_id_sigma	0.007	0.005	0.0013	0.017	10	9	1.31	0.0015	0.0013

Process results for plotting¶

In [20]:

Copied!





# Add trialwise parameters to idata
mSimCenteredSampled = mSimCentered.add_likelihood_parameters_to_datatree(
    mSimCenteredSampled
)

plot_df_mcmc = process_idata_for_plotting(
    idata=mSimCenteredSampled, parameter_matrix=parameter_matrix, model="ddm"
)

plot_df_vi = process_idata_for_plotting(
    idata=mSimCenteredVIObject, parameter_matrix=parameter_matrix, model="ddm"
)
# Add trialwise parameters to idata
mSimCenteredSampled = mSimCentered.add_likelihood_parameters_to_datatree(
    mSimCenteredSampled
)

plot_df_mcmc = process_idata_for_plotting(
    idata=mSimCenteredSampled, parameter_matrix=parameter_matrix, model="ddm"
)

plot_df_vi = process_idata_for_plotting(
    idata=mSimCenteredVIObject, parameter_matrix=parameter_matrix, model="ddm"
)

Plotting¶

In [21]:

Copied!

hssm.plotting.plot_predictive(mSimCentered, col="participant_id", col_wrap=5)
hssm.plotting.plot_predictive(mSimCentered, col="participant_id", col_wrap=5)

Out[21]:

<seaborn.axisgrid.FacetGrid at 0x13cf80830>

In [22]:

Copied!





# Suppose your dataframe is called `df`
params = hssm.defaults.default_model_config["ddm"]["list_params"]
fig, axes = plt.subplots(nrows=1, ncols=len(params), figsize=(14, 6), sharey=True)

for ax, par in zip(axes, params):
    # Identify the relevant columns for this parameter
    mean_col = f"{par}_mean"
    lower_col = f"{par}_lower_hdi"
    upper_col = f"{par}_upper_hdi"
    true_col = f"{par}_true"  # if you also want to plot 'true' values

    # Sort if you want the participants in order on the y‐axis
    df_sorted_vi = plot_df_vi.sort_values("participant_id")
    df_sorted_mcmc = plot_df_mcmc.sort_values("participant_id")

    # We'll use the row index (0..N-1) for plotting against x=mean
    yvals = np.array(range(len(df_sorted_vi)))
    yvals_vi = yvals + 0.1
    yvals_mcmc = yvals - 0.1

    # Draw horizontal lines from lower_HDI to upper_HDI
    ax.hlines(
        yvals_vi,
        df_sorted_vi[lower_col],
        df_sorted_vi[upper_col],
        color="blue",
        alpha=0.5,
    )

    ax.hlines(
        yvals_mcmc,
        df_sorted_mcmc[lower_col],
        df_sorted_mcmc[upper_col],
        color="green",
        alpha=0.5,
    )

    # Mark the posterior mean in blue
    ax.plot(
        df_sorted_vi[mean_col],
        yvals_vi,
        "|",
        color="blue",
        label="Mean_vi" if par == "v" else None,
    )
    ax.plot(
        df_sorted_mcmc[mean_col],
        yvals_mcmc,
        "|",
        color="green",
        label="Mean_mcmc" if par == "v" else None,
    )

    # Optionally, if you have true values, plot them as red 'x'
    if true_col in df_sorted_vi.columns:
        ax.plot(
            df_sorted_vi[true_col],
            yvals,
            "x",
            color="red",
            label="True_vi" if par == "v" else None,
        )

    ax.set_title(par)
    ax.set_yticks(yvals)
    ax.set_yticklabels(df_sorted_vi["participant_id"])
    ax.invert_yaxis()  # optional, if you prefer subject_01 at the top
    if par == "v":
        ax.legend()

plt.tight_layout()
plt.show()
# Suppose your dataframe is called `df`
params = hssm.defaults.default_model_config["ddm"]["list_params"]
fig, axes = plt.subplots(nrows=1, ncols=len(params), figsize=(14, 6), sharey=True)

for ax, par in zip(axes, params):
    # Identify the relevant columns for this parameter
    mean_col = f"{par}_mean"
    lower_col = f"{par}_lower_hdi"
    upper_col = f"{par}_upper_hdi"
    true_col = f"{par}_true"  # if you also want to plot 'true' values

    # Sort if you want the participants in order on the y‐axis
    df_sorted_vi = plot_df_vi.sort_values("participant_id")
    df_sorted_mcmc = plot_df_mcmc.sort_values("participant_id")

    # We'll use the row index (0..N-1) for plotting against x=mean
    yvals = np.array(range(len(df_sorted_vi)))
    yvals_vi = yvals + 0.1
    yvals_mcmc = yvals - 0.1

    # Draw horizontal lines from lower_HDI to upper_HDI
    ax.hlines(
        yvals_vi,
        df_sorted_vi[lower_col],
        df_sorted_vi[upper_col],
        color="blue",
        alpha=0.5,
    )

    ax.hlines(
        yvals_mcmc,
        df_sorted_mcmc[lower_col],
        df_sorted_mcmc[upper_col],
        color="green",
        alpha=0.5,
    )

    # Mark the posterior mean in blue
    ax.plot(
        df_sorted_vi[mean_col],
        yvals_vi,
        "|",
        color="blue",
        label="Mean_vi" if par == "v" else None,
    )
    ax.plot(
        df_sorted_mcmc[mean_col],
        yvals_mcmc,
        "|",
        color="green",
        label="Mean_mcmc" if par == "v" else None,
    )

    # Optionally, if you have true values, plot them as red 'x'
    if true_col in df_sorted_vi.columns:
        ax.plot(
            df_sorted_vi[true_col],
            yvals,
            "x",
            color="red",
            label="True_vi" if par == "v" else None,
        )

    ax.set_title(par)
    ax.set_yticks(yvals)
    ax.set_yticklabels(df_sorted_vi["participant_id"])
    ax.invert_yaxis()  # optional, if you prefer subject_01 at the top
    if par == "v":
        ax.legend()

plt.tight_layout()
plt.show()

Here, results match up quite nicely between the two approaches. Overall the posteriors are very similar, and we don't see a strong tendency for VI posteriors to be more peaked than MCMC posteriors. However this is just an example, and we can't deduce a general rule from this observation.

Hierarchical Model: Non-Centered Parameterization¶

To highlight that once we have the trial wise parameters, we can easily generate subject (and/or condition) wise plots of the posteriors, we also show an example using the non-centered parameterization.

The non-centered parameterization is slightly more complex, so manually recomposing parameters can sometimes be a little bit more confusing.

We can let HSSM handle this via Bambi under the hood, and not worry about this complexity.

In [23]:

Copied!





# Generative model
mSimNonCentered = hssm.HSSM(
    data=simDataDDM,
    p_outlier=0.01,
    prior_settings="safe",
    noncentered=True,
    model="ddm",
    loglik_kind="approx_differentiable",
    include=[
        {
            "name": "v",
            "formula": "v ~ 1 + (1|participant_id)",
        },
        {
            "name": "a",
            "formula": "a ~ 1 + (1|participant_id)",
        },
        {
            "name": "z",
            "formula": "z ~ 1 + (1|participant_id)",
        },
        {
            "name": "t",
            "formula": "t ~ 1 + (1|participant_id)",
        },
    ],
)
# Generative model
mSimNonCentered = hssm.HSSM(
    data=simDataDDM,
    p_outlier=0.01,
    prior_settings="safe",
    noncentered=True,
    model="ddm",
    loglik_kind="approx_differentiable",
    include=[
        {
            "name": "v",
            "formula": "v ~ 1 + (1|participant_id)",
        },
        {
            "name": "a",
            "formula": "a ~ 1 + (1|participant_id)",
        },
        {
            "name": "z",
            "formula": "z ~ 1 + (1|participant_id)",
        },
        {
            "name": "t",
            "formula": "t ~ 1 + (1|participant_id)",
        },
    ],
)

Model initialized successfully.

Fit VI¶

In [24]:

Copied!





# VI
vi_idata_NC = mSimNonCentered.vi(
    niter=30000,
    method="advi",
    obj_optimizer=pm.adamax(learning_rate=0.01),
)
with mSimNonCentered.pymc_model:
    mSimNonCenteredVIObject = mSimNonCentered.vi_approx.sample(draws=1000)
# VI
vi_idata_NC = mSimNonCentered.vi(
    niter=30000,
    method="advi",
    obj_optimizer=pm.adamax(learning_rate=0.01),
)
with mSimNonCentered.pymc_model:
    mSimNonCenteredVIObject = mSimNonCentered.vi_approx.sample(draws=1000)

Using MCMC starting point defaults.

pytensor/link/jax/linker.py:50: UserWarning: The RandomType SharedVariables [RNG(<Generator(PCG64) at 0x13CEE7AE0>)] will not be used in the compiled JAX graph. Instead a copy will be used.
  warnings.warn(

Finished [100%]: Average Loss = 3,320.7

In [25]:

Copied!





# Loss plot
plt.plot(mSimNonCentered.vi_approx.hist)
plt.xlabel("Iteration")
plt.ylabel("Loss")
plt.title("VI iteration loss")
plt.show()
# Loss plot
plt.plot(mSimNonCentered.vi_approx.hist)
plt.xlabel("Iteration")
plt.ylabel("Loss")
plt.title("VI iteration loss")
plt.show()

Fit MCMC¶

In [26]:

Copied!





# MCMC
mSimNonCenteredSampled = mSimNonCentered.sample(
    sampler="numpyro",
    cores=4,
    chains=4,
    draws=250,
    tune=500,
    mp_ctx="forkserver",
    nuts={"max_tree_depth": 5},
)

mSimNonCentered.sample_posterior_predictive(dt=mSimNonCenteredSampled)
# MCMC
mSimNonCenteredSampled = mSimNonCentered.sample(
    sampler="numpyro",
    cores=4,
    chains=4,
    draws=250,
    tune=500,
    mp_ctx="forkserver",
    nuts={"max_tree_depth": 5},
)

mSimNonCentered.sample_posterior_predictive(dt=mSimNonCenteredSampled)

Using default initvals.

bambi/backend/pymc.py:232: UserWarning: `init='adapt_diag'` is ignored by `nuts_sampler='numpyro'`; the external sampler uses its own initialization.
  idata = pm.sample(
NUTS[numpyro]: [v_Intercept, v_1|participant_id_sigma, v_1|participant_id_offset, a_Intercept, a_1|participant_id_sigma, a_1|participant_id_offset, z_Intercept, z_1|participant_id_sigma, z_1|participant_id_offset, t_Intercept, t_1|participant_id_sigma, t_1|participant_id_offset]

pymc/sampling/jax.py:463: UserWarning: There are not enough devices to run parallel chains: expected 4 but got 1. Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(4)` at the beginning of your program. You can double-check how many devices are available in your system using `jax.local_device_count()`.
  pmap_numpyro = MCMC(

  0%|          | 0/750 [00:00<?, ?it/s]

warmup:   5%|▍         | 37/750 [00:02<00:40, 17.62it/s, 31 steps of size 2.42e-02. acc. prob=0.74]

warmup:  10%|▉         | 74/750 [00:04<00:36, 18.61it/s, 31 steps of size 8.07e-03. acc. prob=0.75]

warmup:  15%|█▍        | 111/750 [00:05<00:33, 18.87it/s, 31 steps of size 1.50e-01. acc. prob=0.77]

warmup:  20%|█▉        | 148/750 [00:07<00:32, 18.68it/s, 31 steps of size 4.00e-01. acc. prob=0.78]

warmup:  25%|██▍       | 185/750 [00:09<00:28, 19.58it/s, 15 steps of size 1.96e-01. acc. prob=0.78]

warmup:  30%|██▉       | 222/750 [00:11<00:25, 20.44it/s, 31 steps of size 1.87e-01. acc. prob=0.78]

warmup:  35%|███▍      | 259/750 [00:12<00:23, 21.01it/s, 31 steps of size 1.07e-01. acc. prob=0.78]

warmup:  39%|███▉      | 296/750 [00:14<00:21, 20.77it/s, 31 steps of size 1.75e-01. acc. prob=0.78]

warmup:  44%|████▍     | 333/750 [00:16<00:19, 20.93it/s, 31 steps of size 1.04e-01. acc. prob=0.78]

warmup:  49%|████▉     | 370/750 [00:18<00:18, 20.71it/s, 15 steps of size 6.53e-02. acc. prob=0.78]

warmup:  54%|█████▍    | 407/750 [00:20<00:17, 19.89it/s, 31 steps of size 1.82e-01. acc. prob=0.78]

warmup:  59%|█████▉    | 444/750 [00:22<00:15, 19.96it/s, 31 steps of size 1.00e-01. acc. prob=0.78]

warmup:  64%|██████▍   | 481/750 [00:24<00:13, 19.83it/s, 31 steps of size 1.68e-01. acc. prob=0.78]

sample:  69%|██████▉   | 518/750 [00:26<00:11, 19.77it/s, 31 steps of size 1.28e-01. acc. prob=0.93]

sample:  74%|███████▍  | 555/750 [00:28<00:10, 19.23it/s, 31 steps of size 1.28e-01. acc. prob=0.92]

sample:  79%|███████▉  | 592/750 [00:30<00:08, 18.83it/s, 31 steps of size 1.28e-01. acc. prob=0.91]

sample:  84%|████████▍ | 629/750 [00:32<00:06, 18.58it/s, 31 steps of size 1.28e-01. acc. prob=0.91]

sample:  89%|████████▉ | 666/750 [00:34<00:04, 18.39it/s, 31 steps of size 1.28e-01. acc. prob=0.90]

sample:  94%|█████████▎| 703/750 [00:36<00:02, 18.27it/s, 31 steps of size 1.28e-01. acc. prob=0.90]

sample:  99%|█████████▊| 740/750 [00:38<00:00, 18.11it/s, 31 steps of size 1.28e-01. acc. prob=0.90]

sample: 100%|██████████| 750/750 [00:39<00:00, 19.22it/s, 31 steps of size 1.28e-01. acc. prob=0.90]

  0%|          | 0/750 [00:00<?, ?it/s]

warmup:   5%|▍         | 37/750 [00:01<00:29, 24.25it/s, 31 steps of size 8.84e-03. acc. prob=0.73]

warmup:  10%|▉         | 74/750 [00:03<00:32, 20.70it/s, 31 steps of size 1.07e-02. acc. prob=0.75]

warmup:  15%|█▍        | 111/750 [00:05<00:32, 19.49it/s, 31 steps of size 1.96e-01. acc. prob=0.77]

warmup:  20%|█▉        | 148/750 [00:07<00:31, 18.97it/s, 31 steps of size 4.84e-01. acc. prob=0.78]

warmup:  25%|██▍       | 185/750 [00:09<00:28, 19.73it/s, 31 steps of size 5.78e-02. acc. prob=0.77]

warmup:  30%|██▉       | 222/750 [00:11<00:27, 19.33it/s, 31 steps of size 1.52e-01. acc. prob=0.78]

warmup:  35%|███▍      | 259/750 [00:13<00:25, 19.53it/s, 15 steps of size 3.27e-01. acc. prob=0.78]

warmup:  39%|███▉      | 296/750 [00:14<00:22, 20.33it/s, 31 steps of size 3.38e-01. acc. prob=0.78]

warmup:  44%|████▍     | 333/750 [00:16<00:19, 20.91it/s, 15 steps of size 2.50e-01. acc. prob=0.78]

warmup:  49%|████▉     | 370/750 [00:18<00:17, 21.41it/s, 15 steps of size 1.14e-01. acc. prob=0.78]

warmup:  54%|█████▍    | 407/750 [00:19<00:16, 21.05it/s, 31 steps of size 1.69e-01. acc. prob=0.78]

warmup:  59%|█████▉    | 444/750 [00:21<00:14, 21.23it/s, 31 steps of size 1.57e-01. acc. prob=0.78]

warmup:  64%|██████▍   | 481/750 [00:23<00:12, 21.01it/s, 31 steps of size 1.02e-01. acc. prob=0.78]

sample:  69%|██████▉   | 518/750 [00:25<00:11, 20.26it/s, 31 steps of size 1.11e-01. acc. prob=0.95]

sample:  74%|███████▍  | 555/750 [00:27<00:10, 19.48it/s, 31 steps of size 1.11e-01. acc. prob=0.94]

sample:  79%|███████▉  | 592/750 [00:29<00:08, 18.97it/s, 31 steps of size 1.11e-01. acc. prob=0.93]

sample:  84%|████████▍ | 629/750 [00:31<00:06, 18.61it/s, 31 steps of size 1.11e-01. acc. prob=0.94]

sample:  89%|████████▉ | 666/750 [00:33<00:04, 18.36it/s, 31 steps of size 1.11e-01. acc. prob=0.93]

sample:  94%|█████████▎| 703/750 [00:35<00:02, 18.19it/s, 31 steps of size 1.11e-01. acc. prob=0.93]

sample:  99%|█████████▊| 740/750 [00:37<00:00, 18.13it/s, 31 steps of size 1.11e-01. acc. prob=0.93]

sample: 100%|██████████| 750/750 [00:38<00:00, 19.53it/s, 31 steps of size 1.11e-01. acc. prob=0.93]

  0%|          | 0/750 [00:00<?, ?it/s]

warmup:   5%|▍         | 37/750 [00:01<00:33, 21.11it/s, 31 steps of size 1.14e-02. acc. prob=0.73]

warmup:  10%|▉         | 74/750 [00:03<00:33, 19.90it/s, 31 steps of size 1.84e-02. acc. prob=0.76]

warmup:  15%|█▍        | 111/750 [00:05<00:33, 19.13it/s, 31 steps of size 4.38e-01. acc. prob=0.77]

warmup:  20%|█▉        | 148/750 [00:07<00:31, 18.90it/s, 31 steps of size 7.35e-01. acc. prob=0.78]

warmup:  25%|██▍       | 185/750 [00:09<00:28, 19.52it/s, 31 steps of size 1.36e-01. acc. prob=0.78]

warmup:  30%|██▉       | 222/750 [00:11<00:26, 19.69it/s, 15 steps of size 2.84e-01. acc. prob=0.78]

warmup:  35%|███▍      | 259/750 [00:13<00:24, 20.40it/s, 31 steps of size 6.68e-02. acc. prob=0.78]

warmup:  39%|███▉      | 296/750 [00:15<00:23, 19.59it/s, 31 steps of size 1.97e-01. acc. prob=0.78]

warmup:  44%|████▍     | 333/750 [00:17<00:21, 19.01it/s, 31 steps of size 1.60e-01. acc. prob=0.78]

warmup:  49%|████▉     | 370/750 [00:19<00:20, 18.70it/s, 31 steps of size 1.07e-01. acc. prob=0.78]

warmup:  54%|█████▍    | 407/750 [00:21<00:18, 18.43it/s, 14 steps of size 5.65e-02. acc. prob=0.78]

warmup:  59%|█████▉    | 444/750 [00:23<00:16, 18.13it/s, 31 steps of size 1.38e-01. acc. prob=0.78]

warmup:  64%|██████▍   | 481/750 [00:25<00:14, 18.90it/s, 31 steps of size 6.40e-02. acc. prob=0.78]

sample:  69%|██████▉   | 518/750 [00:27<00:12, 18.56it/s, 31 steps of size 1.40e-01. acc. prob=0.89]

sample:  74%|███████▍  | 555/750 [00:29<00:10, 18.25it/s, 31 steps of size 1.40e-01. acc. prob=0.88]

sample:  79%|███████▉  | 592/750 [00:31<00:08, 18.06it/s, 31 steps of size 1.40e-01. acc. prob=0.90]

sample:  84%|████████▍ | 629/750 [00:33<00:06, 17.97it/s, 31 steps of size 1.40e-01. acc. prob=0.90]

sample:  89%|████████▉ | 666/750 [00:35<00:04, 17.94it/s, 31 steps of size 1.40e-01. acc. prob=0.91]

sample:  94%|█████████▎| 703/750 [00:37<00:02, 17.90it/s, 31 steps of size 1.40e-01. acc. prob=0.91]

sample:  99%|█████████▊| 740/750 [00:39<00:00, 17.87it/s, 31 steps of size 1.40e-01. acc. prob=0.91]

sample: 100%|██████████| 750/750 [00:40<00:00, 18.61it/s, 31 steps of size 1.40e-01. acc. prob=0.91]

  0%|          | 0/750 [00:00<?, ?it/s]

warmup:   5%|▍         | 37/750 [00:01<00:33, 21.05it/s, 31 steps of size 1.39e-02. acc. prob=0.73]

warmup:  10%|▉         | 74/750 [00:03<00:34, 19.43it/s, 31 steps of size 1.31e-02. acc. prob=0.76]

warmup:  15%|█▍        | 111/750 [00:05<00:33, 18.92it/s, 31 steps of size 4.18e-01. acc. prob=0.77]

warmup:  20%|█▉        | 148/750 [00:07<00:31, 18.98it/s, 31 steps of size 3.28e-01. acc. prob=0.78]

warmup:  25%|██▍       | 185/750 [00:09<00:27, 20.37it/s, 31 steps of size 2.54e-01. acc. prob=0.78]

warmup:  30%|██▉       | 222/750 [00:10<00:25, 21.10it/s, 31 steps of size 1.57e-01. acc. prob=0.78]

warmup:  35%|███▍      | 259/750 [00:12<00:23, 21.24it/s, 31 steps of size 1.62e-01. acc. prob=0.78]

warmup:  39%|███▉      | 296/750 [00:14<00:21, 20.88it/s, 31 steps of size 1.21e-01. acc. prob=0.78]

warmup:  44%|████▍     | 333/750 [00:16<00:20, 20.84it/s, 31 steps of size 9.63e-02. acc. prob=0.78]

warmup:  49%|████▉     | 370/750 [00:18<00:18, 20.76it/s, 15 steps of size 2.28e-01. acc. prob=0.78]

warmup:  54%|█████▍    | 407/750 [00:19<00:16, 20.64it/s, 15 steps of size 2.02e-01. acc. prob=0.78]

warmup:  59%|█████▉    | 444/750 [00:21<00:14, 20.59it/s, 15 steps of size 1.37e-01. acc. prob=0.78]

warmup:  64%|██████▍   | 481/750 [00:23<00:13, 20.53it/s, 31 steps of size 2.01e-01. acc. prob=0.78]

sample:  69%|██████▉   | 518/750 [00:25<00:11, 19.86it/s, 31 steps of size 8.94e-02. acc. prob=0.95]

sample:  74%|███████▍  | 555/750 [00:27<00:10, 19.19it/s, 31 steps of size 8.94e-02. acc. prob=0.95]

sample:  79%|███████▉  | 592/750 [00:29<00:08, 18.48it/s, 31 steps of size 8.94e-02. acc. prob=0.95]

sample:  84%|████████▍ | 629/750 [00:31<00:06, 18.22it/s, 31 steps of size 8.94e-02. acc. prob=0.95]

sample:  89%|████████▉ | 666/750 [00:33<00:04, 18.11it/s, 31 steps of size 8.94e-02. acc. prob=0.94]

sample:  94%|█████████▎| 703/750 [00:36<00:02, 18.02it/s, 31 steps of size 8.94e-02. acc. prob=0.94]

sample:  99%|█████████▊| 740/750 [00:38<00:00, 17.93it/s, 31 steps of size 8.94e-02. acc. prob=0.95]

sample: 100%|██████████| 750/750 [00:38<00:00, 19.40it/s, 31 steps of size 8.94e-02. acc. prob=0.94]

The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details

The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run ddm_RV_rv{"(),(),(),(),()->(2)"}'s perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

Post-process VI / MCMC results¶

In [27]:

Copied!





mSimNonCenteredSampled = mSimNonCentered.add_likelihood_parameters_to_datatree(
    mSimNonCenteredSampled
)

plot_df_vi_nc = process_idata_for_plotting(
    idata=mSimNonCenteredVIObject, parameter_matrix=parameter_matrix, model="ddm"
)

plot_df_mcmc_nc = process_idata_for_plotting(
    idata=mSimNonCenteredSampled, parameter_matrix=parameter_matrix, model="ddm"
)
mSimNonCenteredSampled = mSimNonCentered.add_likelihood_parameters_to_datatree(
    mSimNonCenteredSampled
)

plot_df_vi_nc = process_idata_for_plotting(
    idata=mSimNonCenteredVIObject, parameter_matrix=parameter_matrix, model="ddm"
)

plot_df_mcmc_nc = process_idata_for_plotting(
    idata=mSimNonCenteredSampled, parameter_matrix=parameter_matrix, model="ddm"
)

Plotting¶

In [28]:

Copied!

hssm.plotting.plot_predictive(mSimNonCentered, col="participant_id", col_wrap=5)
hssm.plotting.plot_predictive(mSimNonCentered, col="participant_id", col_wrap=5)

Out[28]:

<seaborn.axisgrid.FacetGrid at 0x13b206990>

In [29]:

Copied!





# Suppose your dataframe is called `df`
params = hssm.defaults.default_model_config["ddm"]["list_params"]
fig, axes = plt.subplots(nrows=1, ncols=len(params), figsize=(14, 6), sharey=True)

for ax, par in zip(axes, params):
    # Identify the relevant columns for this parameter
    mean_col = f"{par}_mean"
    lower_col = f"{par}_lower_hdi"
    upper_col = f"{par}_upper_hdi"
    true_col = f"{par}_true"  # if you also want to plot 'true' values

    # Sort if you want the participants in order on the y‐axis
    df_sorted_vi_nc = plot_df_vi_nc.sort_values("participant_id")
    df_sorted_mcmc_nc = plot_df_mcmc_nc.sort_values("participant_id")

    # We'll use the row index (0..N-1) for plotting against x=mean
    yvals = np.array(range(len(df_sorted_vi_nc)))
    yvals_vi = yvals + 0.1
    yvals_mcmc = yvals - 0.1

    # Draw horizontal lines from lower_HDI to upper_HDI
    ax.hlines(
        yvals_vi,
        df_sorted_vi_nc[lower_col],
        df_sorted_vi_nc[upper_col],
        color="blue",
        alpha=0.5,
    )

    ax.hlines(
        yvals_mcmc,
        df_sorted_mcmc_nc[lower_col],
        df_sorted_mcmc_nc[upper_col],
        color="green",
        alpha=0.5,
    )

    # Mark the posterior mean in blue
    ax.plot(
        df_sorted_vi_nc[mean_col],
        yvals,
        "|",
        color="blue",
        label="Mean_vi" if par == "v" else None,
    )
    ax.plot(
        df_sorted_mcmc_nc[mean_col],
        yvals_mcmc,
        "|",
        color="green",
        label="Mean_mcmc" if par == "v" else None,
    )

    # Optionally, if you have true values, plot them as red 'x'
    if true_col in df_sorted_vi_nc.columns:
        ax.plot(
            df_sorted_vi_nc[true_col],
            yvals,
            "x",
            color="red",
            label="True_vi" if par == "v" else None,
        )

    ax.set_title(par)
    ax.set_yticks(yvals)
    ax.set_yticklabels(df_sorted_vi_nc["participant_id"])
    ax.invert_yaxis()  # optional, if you prefer subject_01 at the top
    if par == "v":
        ax.legend()

plt.tight_layout()
plt.show()
# Suppose your dataframe is called `df`
params = hssm.defaults.default_model_config["ddm"]["list_params"]
fig, axes = plt.subplots(nrows=1, ncols=len(params), figsize=(14, 6), sharey=True)

for ax, par in zip(axes, params):
    # Identify the relevant columns for this parameter
    mean_col = f"{par}_mean"
    lower_col = f"{par}_lower_hdi"
    upper_col = f"{par}_upper_hdi"
    true_col = f"{par}_true"  # if you also want to plot 'true' values

    # Sort if you want the participants in order on the y‐axis
    df_sorted_vi_nc = plot_df_vi_nc.sort_values("participant_id")
    df_sorted_mcmc_nc = plot_df_mcmc_nc.sort_values("participant_id")

    # We'll use the row index (0..N-1) for plotting against x=mean
    yvals = np.array(range(len(df_sorted_vi_nc)))
    yvals_vi = yvals + 0.1
    yvals_mcmc = yvals - 0.1

    # Draw horizontal lines from lower_HDI to upper_HDI
    ax.hlines(
        yvals_vi,
        df_sorted_vi_nc[lower_col],
        df_sorted_vi_nc[upper_col],
        color="blue",
        alpha=0.5,
    )

    ax.hlines(
        yvals_mcmc,
        df_sorted_mcmc_nc[lower_col],
        df_sorted_mcmc_nc[upper_col],
        color="green",
        alpha=0.5,
    )

    # Mark the posterior mean in blue
    ax.plot(
        df_sorted_vi_nc[mean_col],
        yvals,
        "|",
        color="blue",
        label="Mean_vi" if par == "v" else None,
    )
    ax.plot(
        df_sorted_mcmc_nc[mean_col],
        yvals_mcmc,
        "|",
        color="green",
        label="Mean_mcmc" if par == "v" else None,
    )

    # Optionally, if you have true values, plot them as red 'x'
    if true_col in df_sorted_vi_nc.columns:
        ax.plot(
            df_sorted_vi_nc[true_col],
            yvals,
            "x",
            color="red",
            label="True_vi" if par == "v" else None,
        )

    ax.set_title(par)
    ax.set_yticks(yvals)
    ax.set_yticklabels(df_sorted_vi_nc["participant_id"])
    ax.invert_yaxis()  # optional, if you prefer subject_01 at the top
    if par == "v":
        ax.legend()

plt.tight_layout()
plt.show()

Overall posteriors look very similar between the two approaches again. However, we can see that the MCMC posteriors are more variable than the VI posteriors in this case. This illustrates a case where VI posteriors are too confident, and bayesian t-tests on parameter differences may yield different results when using the MCMC vs. the VI posteriors. From a practical perspective, it is sometimes simply infeasible to run MCMC, while VI is still computationally tractable. In such cases we have to take what we can get...

This is the pattern to remember from the whole case study: VI's failure mode is overconfidence, and it is parameterization- and model-dependent — there is no general rule, which is why the decision guide above recommends verifying VI against a short MCMC run for any model you intend to rely on.

Appendix: driving PyMC's VI API directly¶

Here we illustrate how to use our attached pymc_model to make use of the object oriented API for variational inference. This allows us a few extra affordances.

Let's define a few helper functions first.

In [30]:

Copied!





import warnings

import xarray as xr
from matplotlib import gridspec
from pymc.blocking import DictToArrayBijection, RaveledVars


def tracker_to_idata(tracker, model):
    """Turn a tracker object into an DataTree object."""
    tracker_groups = list(tracker.whatchdict.keys())
    # n_steps = len(tracker[tracker_groups[0]])
    stacked_results = {
        tracker_group: {
            key: np.stack([d[key] for d in tracker[tracker_group]])
            for key in tracker[tracker_group][0]
        }
        for tracker_group in tracker_groups
    }

    # coords = {"vi_step": np.arange(n_steps)} | {
    #     k: np.array(v) for k, v in model.coords.items()
    # }
    var_to_dims = {
        var.name: ("vi_step", *(model.named_vars_to_dims.get(var.name, ())))
        for var in model.continuous_value_vars
    }
    datasets = {
        key: xr.Dataset(
            {
                var: (var_to_dims[var], stacked_results[key][var])
                for var in stacked_results[key].keys()
            }
        )
        for key in tracker_groups
    }

    with warnings.catch_warnings(action="ignore"):
        return xr.DataTree.from_dict(datasets)


def untransform_params(idata, model):
    """Bring transformed parmater back to their original scale."""
    suffixes = ["_interval__", "_log__"]

    def remove_suffixes(word, suffixes):
        for suffix in suffixes:
            if word.endswith(suffix):
                return word[: -len(suffix)]
        return word

    free_rv_names = [rv_.name for rv_ in model.free_RVs]
    transformed_vars = list(idata["mean"].ds.data_vars.keys())
    collect_untransformed_vars = []
    collect_untransformed_xarray_datasets = []

    for var_ in transformed_vars:
        var_untrans = remove_suffixes(var_, suffixes=suffixes)
        if var_untrans in free_rv_names:
            rv = model.free_RVs[free_rv_names.index(var_untrans)]
            if model.rvs_to_transforms[rv] is not None:
                untransformed_var = (
                    model.rvs_to_transforms[rv]
                    .backward(idata["mean"][var_].values, *rv.owner.inputs)
                    .eval()
                )
                collect_untransformed_vars.append(var_)
                collect_untransformed_xarray_datasets.append(
                    xr.Dataset(
                        data_vars={var_untrans: (("vi_step"), untransformed_var)}
                    )
                )

    return xr.merge(
        [idata["mean"].ds] + collect_untransformed_xarray_datasets
    ).drop_vars(collect_untransformed_vars)


def plot_vi_traces(idata):
    """Plot parameter history of the optimization alogrithm."""
    if not isinstance(idata, xr.DataTree):
        raise ValueError("idata must be a DataTree object")
    if "loss" not in idata:
        raise ValueError("DataTree object must contain a 'loss' group")
    if "mean_untransformed" not in idata:
        print(
            "Using transformed variables because 'mean_untransformed' group not found"
        )
        data_vars = list(idata["mean"].ds.data_vars.keys())
    else:
        data_vars = list(idata["mean_untransformed"].ds.data_vars.keys())

    fig = plt.figure(figsize=(8, 1.5 * len(data_vars)))
    gs = gridspec.GridSpec(
        len(data_vars) // 2 + 2
        if (len(data_vars) % 2) == 0
        else (len(data_vars) // 2) + 3,
        2,
    )

    for i, var_ in enumerate(data_vars):
        ax_tmp = fig.add_subplot(gs[i // 2, i % 2])
        idata["mean_untransformed"][var_].plot(ax=ax_tmp)
        ax_tmp.set_title(var_)

    last_ax = fig.add_subplot(gs[-2:, :])
    idata["loss"].loss.plot(ax=last_ax)
    gs.tight_layout(fig)
    return fig
import warnings

import xarray as xr
from matplotlib import gridspec
from pymc.blocking import DictToArrayBijection, RaveledVars


def tracker_to_idata(tracker, model):
    """Turn a tracker object into an DataTree object."""
    tracker_groups = list(tracker.whatchdict.keys())
    # n_steps = len(tracker[tracker_groups[0]])
    stacked_results = {
        tracker_group: {
            key: np.stack([d[key] for d in tracker[tracker_group]])
            for key in tracker[tracker_group][0]
        }
        for tracker_group in tracker_groups
    }

    # coords = {"vi_step": np.arange(n_steps)} | {
    #     k: np.array(v) for k, v in model.coords.items()
    # }
    var_to_dims = {
        var.name: ("vi_step", *(model.named_vars_to_dims.get(var.name, ())))
        for var in model.continuous_value_vars
    }
    datasets = {
        key: xr.Dataset(
            {
                var: (var_to_dims[var], stacked_results[key][var])
                for var in stacked_results[key].keys()
            }
        )
        for key in tracker_groups
    }

    with warnings.catch_warnings(action="ignore"):
        return xr.DataTree.from_dict(datasets)


def untransform_params(idata, model):
    """Bring transformed parmater back to their original scale."""
    suffixes = ["_interval__", "_log__"]

    def remove_suffixes(word, suffixes):
        for suffix in suffixes:
            if word.endswith(suffix):
                return word[: -len(suffix)]
        return word

    free_rv_names = [rv_.name for rv_ in model.free_RVs]
    transformed_vars = list(idata["mean"].ds.data_vars.keys())
    collect_untransformed_vars = []
    collect_untransformed_xarray_datasets = []

    for var_ in transformed_vars:
        var_untrans = remove_suffixes(var_, suffixes=suffixes)
        if var_untrans in free_rv_names:
            rv = model.free_RVs[free_rv_names.index(var_untrans)]
            if model.rvs_to_transforms[rv] is not None:
                untransformed_var = (
                    model.rvs_to_transforms[rv]
                    .backward(idata["mean"][var_].values, *rv.owner.inputs)
                    .eval()
                )
                collect_untransformed_vars.append(var_)
                collect_untransformed_xarray_datasets.append(
                    xr.Dataset(
                        data_vars={var_untrans: (("vi_step"), untransformed_var)}
                    )
                )

    return xr.merge(
        [idata["mean"].ds] + collect_untransformed_xarray_datasets
    ).drop_vars(collect_untransformed_vars)


def plot_vi_traces(idata):
    """Plot parameter history of the optimization alogrithm."""
    if not isinstance(idata, xr.DataTree):
        raise ValueError("idata must be a DataTree object")
    if "loss" not in idata:
        raise ValueError("DataTree object must contain a 'loss' group")
    if "mean_untransformed" not in idata:
        print(
            "Using transformed variables because 'mean_untransformed' group not found"
        )
        data_vars = list(idata["mean"].ds.data_vars.keys())
    else:
        data_vars = list(idata["mean_untransformed"].ds.data_vars.keys())

    fig = plt.figure(figsize=(8, 1.5 * len(data_vars)))
    gs = gridspec.GridSpec(
        len(data_vars) // 2 + 2
        if (len(data_vars) % 2) == 0
        else (len(data_vars) // 2) + 3,
        2,
    )

    for i, var_ in enumerate(data_vars):
        ax_tmp = fig.add_subplot(gs[i // 2, i % 2])
        idata["mean_untransformed"][var_].plot(ax=ax_tmp)
        ax_tmp.set_title(var_)

    last_ax = fig.add_subplot(gs[-2:, :])
    idata["loss"].loss.plot(ax=last_ax)
    gs.tight_layout(fig)
    return fig

In [31]:

Copied!

# Define the ADVI runner
with cav_model.pymc_model:
    advi = pm.ADVI()
# Define the ADVI runner
with cav_model.pymc_model:
    advi = pm.ADVI()

In [32]:

Copied!





# Set up starting point
start = cav_model.pymc_model.initial_point()
vars_dict = {var.name: var for var in cav_model.pymc_model.continuous_value_vars}
x0 = DictToArrayBijection.map(
    {var_name: value for var_name, value in start.items() if var_name in vars_dict}
)

# Define quantities to track
tracker = pm.variational.callbacks.Tracker(
    mean=lambda: DictToArrayBijection.rmap(
        RaveledVars(advi.approx.mean.eval(), x0.point_map_info), start
    ),  # callable that returns mean
    std=lambda: DictToArrayBijection.rmap(
        RaveledVars(advi.approx.std.eval(), x0.point_map_info), start
    ),  # callable that returns std
)
# Set up starting point
start = cav_model.pymc_model.initial_point()
vars_dict = {var.name: var for var in cav_model.pymc_model.continuous_value_vars}
x0 = DictToArrayBijection.map(
    {var_name: value for var_name, value in start.items() if var_name in vars_dict}
)

# Define quantities to track
tracker = pm.variational.callbacks.Tracker(
    mean=lambda: DictToArrayBijection.rmap(
        RaveledVars(advi.approx.mean.eval(), x0.point_map_info), start
    ),  # callable that returns mean
    std=lambda: DictToArrayBijection.rmap(
        RaveledVars(advi.approx.std.eval(), x0.point_map_info), start
    ),  # callable that returns std
)

In [33]:

Copied!





# Run VI fit

approx = advi.fit(n=30000, callbacks=[tracker])
with cav_model.pymc_model:
    vi_posterior_samples = approx.sample(1000)
    vi_posterior_samples.posterior = vi_posterior_samples.posterior.ds.drop_vars(
        "v_mean"
    )
# Run VI fit

approx = advi.fit(n=30000, callbacks=[tracker])
with cav_model.pymc_model:
    vi_posterior_samples = approx.sample(1000)
    vi_posterior_samples.posterior = vi_posterior_samples.posterior.ds.drop_vars(
        "v_mean"
    )

pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(
pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run LANLogpVJPOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

Finished [100%]: Average Loss = 6,037.3

In [34]:

Copied!





from copy import deepcopy

# Convert tracked quantities to idata
result = tracker_to_idata(tracker, cav_model.pymc_model)

# Add untransformed parameters
result["mean_untransformed"] = untransform_params(
    deepcopy(result), cav_model.pymc_model
)
# Add loss group
result["loss"] = xr.Dataset(data_vars={"loss": ("vi_step", np.array(approx.hist))})
from copy import deepcopy

# Convert tracked quantities to idata
result = tracker_to_idata(tracker, cav_model.pymc_model)

# Add untransformed parameters
result["mean_untransformed"] = untransform_params(
    deepcopy(result), cav_model.pymc_model
)
# Add loss group
result["loss"] = xr.Dataset(data_vars={"loss": ("vi_step", np.array(approx.hist))})

A quick look at our result DataTree object, to understand what happened here.

We now have two additional groups:

mean_untransformed which holds parameter values in the orignal space (instead of the parameters over which the optimization operates, which will always live in an unconstrained space)
loss which holds our training history

Plot Results¶

We can plot the parameter trajectories (histories) over optimization steps, with our little helper function plot_vi_traces().

NOTE:

This is a random run, and we did not thoroughly check if the number of steps we allowed the optimizer were indeed enough to converge.

In [36]:

Copied!

fig = plot_vi_traces(result)
fig = plot_vi_traces(result)

In [37]:

Copied!





# Plot MCMC [nuts]
plot_matrix = az.plot_pair(
    cav_model.traces,
    visuals={
        "scatter": dict(alpha=0.01, color="blue"),
        "dist": dict(alpha=0.6, color="blue"),
    },
)
# Plot VI via .vi() [fullrank_advi]
az.plot_pair(
    cav_model.vi_idata,
    plot_matrix=plot_matrix,
    visuals={
        "scatter": dict(alpha=0.04, color="red"),
        "dist": dict(alpha=0.8, color="red"),
    },
)
# Plot VI via pymc interface [advi]
# (We need to make sure the variables are in correct order)
az.plot_pair(
    vi_posterior_samples.posterior[list(cav_model.traces.posterior.data_vars)],
    plot_matrix=plot_matrix,
    visuals={
        "scatter": dict(alpha=0.04, color="green"),
        "dist": dict(alpha=0.8, color="green"),
    },
)
# Plot MCMC [nuts]
plot_matrix = az.plot_pair(
    cav_model.traces,
    visuals={
        "scatter": dict(alpha=0.01, color="blue"),
        "dist": dict(alpha=0.6, color="blue"),
    },
)
# Plot VI via .vi() [fullrank_advi]
az.plot_pair(
    cav_model.vi_idata,
    plot_matrix=plot_matrix,
    visuals={
        "scatter": dict(alpha=0.04, color="red"),
        "dist": dict(alpha=0.8, color="red"),
    },
)
# Plot VI via pymc interface [advi]
# (We need to make sure the variables are in correct order)
az.plot_pair(
    vi_posterior_samples.posterior[list(cav_model.traces.posterior.data_vars)],
    plot_matrix=plot_matrix,
    visuals={
        "scatter": dict(alpha=0.04, color="green"),
        "dist": dict(alpha=0.8, color="green"),
    },
)

Out[37]:

<arviz_plots.plot_matrix.PlotMatrix at 0x15e1b74d0>

NOTE:

It is expected that the posterior of our last run looks a little worse. We chose to run the advi algorithm, which implies only an isotropic Gaussian approximation to the posterior, so we expect to miss the posterior covariances which we pick up via fullrank_advi as well as MCMC.

Run variational inference¶

Load Data and Specify model¶

Inference¶

Run MCMC¶

Run VI¶

Inspect Outputs¶

.vi_idata¶

.vi_approx¶

Contrast outputs between MCMC and VI¶

When to use VI — and when to stick with MCMC¶

Case study: VI on a hierarchical DDM¶

Utilities¶

Simulate Dataset¶

Hierarchical Model: Centered Parameterization¶

Fit VI¶

Fit MCMC¶

Process results for plotting¶

Plotting¶

Hierarchical Model: Non-Centered Parameterization¶

Fit VI¶

Fit MCMC¶

Post-process VI / MCMC results¶

Plotting¶

Appendix: driving PyMC's VI API directly¶

Plot Results¶

Further Reading¶

`.vi_idata`¶

`.vi_approx`¶