Other NUTS Samplers#

In this notebook we show how to fit a CLV model with other NUTS samplers. These alternative samplers can be significantly faster and also sample on the GPU.

Note

You need to install these packages in your Python environment.

Tip

You can pass the exact same nuts_sampler argument to the MMM models.

Tip

GPU support only works with select samplers in PyMC that use the JAX backend. These samplers include numpyro, blackjax, and nutpie.

Make sure the GPU is registered, follow the instructions here.

For the purpose of illustration, we will use the same data and model as in the other CLV notebooks.

import arviz as az
import arviz_plots as azp
import matplotlib.pyplot as plt
import pandas as pd

from pymc_marketing import clv

az.style.use("arviz-darkgrid")
plt.rcParams["figure.figsize"] = [12, 7]
plt.rcParams["figure.dpi"] = 100
plt.rcParams["figure.facecolor"] = "white"

%load_ext autoreload
%autoreload 2
%config InlineBackend.figure_format = "retina"

df = pd.read_csv(
    "https://raw.githubusercontent.com/pymc-labs/pymc-marketing/main/data/clv_quickstart.csv"
)
df["customer_id"] = range(len(df))

We can pass the keyword argument nuts_sampler to the fit method of the CLV model to specify the NUTS sampler to use. In addition, we can pass additional keyword arguments which will be passed to the pymc.sample method via the model builder layer. For example, we can use the numpyro sampler as:

sampler_kwargs = {
    "draws": 2_000,
    "target_accept": 0.9,
    "chains": 5,
    "random_seed": 42,
}

model = clv.BetaGeoModel(data=df)
idata_numpyro = model.fit(nuts_sampler="numpyro", **sampler_kwargs)

NUTS[numpyro]: [alpha, phi_dropout, kappa_dropout, r]

Similarly, we can use the blackjax sampler as:

idata_blackjax = model.fit(nuts_sampler="blackjax", **sampler_kwargs)

NUTS[blackjax]: [alpha, phi_dropout, kappa_dropout, r]

Running window adaptation
100.00% [1000/1000 00:00<?]
100.00% [1000/1000 00:00<?]
100.00% [1000/1000 00:00<?]
100.00% [1000/1000 00:00<?]
100.00% [1000/1000 00:00<?]
100.00% [2000/2000 00:00<?]
100.00% [2000/2000 00:00<?]
100.00% [2000/2000 00:00<?]
100.00% [2000/2000 00:00<?]
100.00% [2000/2000 00:00<?]

Finally, we can use the nutpie which is a Rust implementation of NUTS.

idata_nutpie = model.fit(nuts_sampler="nutpie", **sampler_kwargs)

NUTS[nutpie]: [alpha, phi_dropout, kappa_dropout, r]







The results from the samplers are almost identical:








Hide code cell source





sampler_labels = ["blackjax", "nutpie", "numpyro"]
data_dict = {
    "blackjax": idata_blackjax,
    "nutpie": idata_nutpie,
    "numpyro": idata_numpyro,
}

pc = azp.plot_dist(
    data_dict,
    var_names=["a", "b", "alpha", "r"],
    col_wrap=2,
    visuals={
        "credible_interval": False,
        "point_estimate": False,
        "point_estimate_text": False,
    },
)
fig = pc.viz["/"]["figure"].values.item()
fig.set_size_inches(12, 8)
for ax in fig.axes:
    for line, label in zip(ax.lines, sampler_labels, strict=False):
        line.set_label(label)
    ax.legend(fontsize=8)
fig.suptitle(
    "Posterior Distributions of model parameters",
    fontsize=18,
    fontweight="bold",
    y=1.05,
);










../../_images/550e089ae2e6f2fc9146597c1253d18f387699f6ac67b8399c21249936839a78.png










%load_ext watermark
%watermark -n -u -v -iv -w -p pymc_marketing,blackjax,numpyro,nutpie,pymc









Last updated: Mon, 29 Jun 2026

Python implementation: CPython
Python version       : 3.13.2
IPython version      : 9.14.0

pymc_marketing: 1.0.0.dev0
blackjax      : 1.5
numpyro       : 0.21.0
nutpie        : 0.16.10
pymc          : 6.0.1

arviz         : 1.2.0
arviz_plots   : 1.2.0
matplotlib    : 3.10.9
numpy         : 2.4.6
pandas        : 2.3.3
pymc          : 6.0.1
pymc_extras   : 0.12.2.dev1+gee8cc37df
pymc_marketing: 1.0.0.dev0

Watermark: 2.6.0