Cannot Specify Other Samplers

[ColtAllen]

When attempting to use a sampler other than the default (NUTS), a ValueError is raised because the model variables are being assigned twice. The following example is for the ParetoNBDModel, but it's the same deal with BetaGeoModel:

import pandas as pd
from pymc_marketing.clv.models import ParetoNBDModel

# load data
url_cdnow_rfm = "https://raw.githubusercontent.com/pymc-labs/pymc-marketing/main/datasets/clv_quickstart.csv"
df = pd.read_csv(url_cdnow_rfm)
df['customer_id'] = df.index

# init and fit model
pnbd = ParetoNBDModel(df)
pnbd.build_model()

with pnbd.model:

    sample_config = {
        "step": pm.Slice(), 
        "draws": 2000, 
        "tune": 1000
    }

    pnbd.fit(**sample_config)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[51], line 18
     10 with pnbd.model:
     12     sample_config = {
     13         "step": pm.Slice(), 
     14         "draws": 2000, 
     15         "tune": 1000
     16     }
---> 18     pnbd.fit(**sample_config)

File ~/Projects/pymc-marketing/pymc_marketing/clv/models/pareto_nbd.py:319, in ParetoNBDModel.fit(self, fit_method, **kwargs)
    313 with warnings.catch_warnings():
    314     warnings.filterwarnings(
    315         message="Optimization Warning: The Op hyp2f1 does not provide a C implementation. As well as being potentially slow, this also disables loop fusion.",
    316         action="ignore",
    317         category=UserWarning,
    318     )
--> 319     super().fit(fit_method, **kwargs)

File ~/Projects/pymc-marketing/pymc_marketing/clv/models/basic.py:47, in CLVModel.fit(self, fit_method, **kwargs)
     30 def fit(
     31     self,
     32     fit_method: str = "mcmc",
     33     **kwargs,
     34 ) -> az.InferenceData:
     35     """Infer model posterior
     36 
     37     Parameters
   (...)
     44         Other keyword arguments passed to the underlying PyMC routines
     45     """
---> 47     self.build_model()
     49     if fit_method == "mcmc":
     50         self._fit_mcmc(**kwargs)

File ~/Projects/pymc-marketing/pymc_marketing/clv/models/pareto_nbd.py:229, in ParetoNBDModel.build_model(self)
    224 def build_model(
    225     self,
    226 ) -> None:
    227     with pm.Model(coords=self.coords) as self.model:
    228         # purchase rate priors
--> 229         r = self.model.register_rv(self.r_prior, name="r")
    230         alpha = self.model.register_rv(self.alpha_prior, name="alpha")
    232         # churn priors

File ~/miniconda3/envs/pymc-dev/lib/python3.8/site-packages/pymc/model.py:1334, in Model.register_rv(self, rv_var, name, observed, total_size, dims, transform, initval)
   1332     self.free_RVs.append(rv_var)
   1333     self.create_value_var(rv_var, transform)
-> 1334     self.add_named_variable(rv_var, dims)
   1335     self.set_initval(rv_var, initval)
   1336 else:

File ~/miniconda3/envs/pymc-dev/lib/python3.8/site-packages/pymc/model.py:1552, in Model.add_named_variable(self, var, dims)
   1550     raise ValueError("Variable is unnamed.")
   1551 if self.named_vars.tree_contains(var.name):
-> 1552     raise ValueError(f"Variable name {var.name} already exists.")
   1554 if dims is not None:
   1555     if isinstance(dims, str):

ValueError: Variable name r already exists.

It's also worth noting that external samplers like numpyro cannot be used at all the way the base CLVModel is currently written.

[ricardoV94]

CC @michaelraczycki

[michaelraczycki]

Choosing other step methods wasn't something I was aware of on the earlier stages, and as I see now, the way step functions are defined in pymc is requiring them to be used with the context on stack, which prevents it from being passed by itself. I'll create a PR where I change sample_model() to include: with self.model: sampler_args = {**self.sampler_config, **kwargs} if 'step' in sampler_args: sampler_args['step'] = sampler_args['step']() idata = pm.sample(**sampler_args) also, to properly use the ModelBuilder fit, the sampler config should be passed as a parameter, or additional kwargs can be passed to fit() itself, but it shouldn't be called by adding one more context on the stack, as the fit in basic.py creates it under the hood.

[michaelraczycki]

@ricardoV94 is it something that should be always allowed for all models (therefore is worth being included on the model builder level), or will be used either only in the CLV context or even only Pareto?

[BBDS-Colt]

(This is @ColtAllen, responding from my work account).

I can't speak for the other models, but before @ricardoV94 made some pytensor enhancements, ParetoNBDModel could only be used with gradient-free methods like SliceSampler. I want to see if the new enhancements enable NUTS to converge more quickly than SliceSampler, but am currently blocked from doing so.

The ability to use external samplers across all models would be a great addition I think, as they tend to be much faster. However, JAX samplers cannot currently be used with the ParetoNBDModel because the hyp2f1 Op in the Likelihood isn't supported. I've also installed nutpie, but haven't tested it yet because rust requires a bit more initial setup for a Macbook.

[ricardoV94]

Just to save you trouble @ColtAllen nutpie won't work either because we didn't implement those Ops on the Numba backend

[ColtAllen]

@lucianopaz showed me how to get the slice sampler to work by replacing pnbd.fit(**sampler_config) with pm.sample(**sampler_config) in the above example.

External samplers can be used via a similar approach that can be shared in a tutorial notebook, and is probably the way to go unless there is a lot of user interest in the fit method supporting external samplers (which could be a rather heavy lift).

[ColtAllen]

Closing this because specifying another sampler is easy with the following convention:

with clv_model.model:
    clv_model.idata = pm.sample(step=pm.SAMPLER())

I do think it's worth adding an option for a gradient-free sampler to ParetoNBDModel though, as well as a nutpie example in the docs. I'll get issues created for both.