Creating custom distribution in PyMC
Sheng Long
The old tutorial in the official documentation is meant for PyMC31. The latest version of PyMC uses a different way, but we will nevertheless use the Generalized Poisson distribution as an example.
1 … where as the latest version is 5.xx
The Generalized Poisson distribution, aka the Lagrangian Poisson distribution, models count data, and is flexible enough to handle both overdispersion and underdispersion. It has the following PMF:
\[ p(Y = y | \theta, \lambda) = \frac{\theta(\theta+\lambda y)^{y-1}e^{-\theta-\lambda y}}{y!} \] where \(\theta > 0\) and \(\max(-1, -\frac{\theta}{4}) \leq \lambda \leq 1\)
Using pymc version 5.26.1We will first need to define the log of the probability function: \[ \log(f(y|\theta, \lambda)) = \log(\theta) + (y-1) \log(\theta + \lambda y) - (\theta + \lambda y) - \log(y!) \]
def genpossion_logp(value: TensorVariable, theta: TensorVariable, lam: TensorVariable):
theta_lam_value = theta + lam * value
log_prob = pm.math.log(theta) + (value - 1) * pm.math.log(theta_lam_value) - theta_lam_value - gammaln(value)
log_prob = pt.switch(theta_lam_value <= 0, -pt.inf, log_prob)
# Handle other constraints (value >= 0, mu > 0, etc.)
log_prob = pt.switch(value < 0, -pt.inf, log_prob)
log_prob = pt.switch(theta <= 0, -pt.inf, log_prob)
log_prob = pt.switch(pt.abs(lam) > 1, -pt.inf, log_prob)
log_prob = pt.switch(lam < -theta / 4, -pt.inf, log_prob)
return log_prob Then define a random sample generator:
def genpoisson_rvs(theta, lam, size=None):
if size is not None:
assert size == theta.shape
else:
size = theta.shape
lam = lam[0]
omega = np.exp(-lam)
X = np.full(size, 0)
S = np.exp(-theta)
P = np.copy(S)
for i in range(size[0]):
U = np.random.uniform()
while U > S[i]:
X[i] += 1
C = theta[i] - lam + lam * X[i]
P[i] = omega * C * (1 + lam / C) ** (X[i] - 1) * P[i] / X[i]
S[i] += P[i]
return XIf you’re searching for “PyMC custom distribution” online you might have encountered this class, pymc.DensityDist2. Note however, that this is part of v4 and therefore is also not the intended way to write things.
2 … source here
The correct way to write things is to use the pymc.CustomDist class, with documentat here.
Cell In[22], line 5 y = pm.CustomDist("y", theta, lam, logp=genpossion_logp, random=genpoisson_rvs, observed=observed=np.random.randn(100, 3)) ^ SyntaxError: invalid syntax
--------------------------------------------------------------------------- TypeError Traceback (most recent call last) File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pytensor\link\utils.py:197, in streamline.<locals>.streamline_default_f() 194 for thunk, node, old_storage in zip( 195 thunks, order, post_thunk_old_storage, strict=False 196 ): --> 197 thunk() 198 for old_s in old_storage: File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pytensor\graph\op.py:545, in Op.make_py_thunk.<locals>.rval(p, i, o, n, cm) 537 @is_thunk_type 538 def rval( 539 p=p, (...) 543 cm=node_compute_map, 544 ): --> 545 r = p(n, [x[0] for x in i], o) 546 for entry in cm: File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pytensor\tensor\random\op.py:428, in RandomVariable.perform(self, node, inputs, outputs) 426 outputs[0][0] = rng 427 outputs[1][0] = np.asarray( --> 428 self.rng_fn(rng, *args, None if size is None else tuple(size)), 429 dtype=self.dtype, 430 ) File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pymc\distributions\custom.py:87, in CustomDistRV.rng_fn(cls, rng, *args) 86 size = args.pop(-1) ---> 87 return cls._random_fn(*args, rng=rng, size=size) TypeError: genpoisson_rvs() got an unexpected keyword argument 'rng' During handling of the above exception, another exception occurred: TypeError Traceback (most recent call last) Cell In[21], line 1 ----> 1 pm.model_to_graphviz(gen_poisson_model) File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pymc\model_graph.py:765, in model_to_graphviz(model, var_names, formatting, save, figsize, dpi, node_formatters, graph_attr, include_dim_lengths) 761 model = modelcontext(model) 762 graph = ModelGraph(model) 763 return make_graph( 764 model.name, --> 765 plates=graph.get_plates(var_names=var_names), 766 edges=graph.edges(var_names=var_names), 767 formatting=formatting, 768 save=save, 769 figsize=figsize, 770 dpi=dpi, 771 graph_attr=graph_attr, 772 node_formatters=node_formatters, 773 create_plate_label=create_plate_label_with_dim_length 774 if include_dim_lengths 775 else create_plate_label_without_dim_length, 776 ) File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pymc\model_graph.py:351, in ModelGraph.get_plates(self, var_names) 344 # TODO: Evaluate all RV shapes at once 345 # This should help find discrepancies, and 346 # avoids unnecessary function compiles for determining labels. 347 dim_lengths: dict[str, int] = { 348 dim_name: fast_eval(value).item() for dim_name, value in self.model.dim_lengths.items() 349 } 350 var_shapes: dict[str, tuple[int, ...]] = { --> 351 var_name: tuple(map(int, fast_eval(self.model[var_name].shape))) 352 for var_name in self.vars_to_plot(var_names) 353 } 355 for var_name in self.vars_to_plot(var_names): 356 shape: tuple[int, ...] = var_shapes[var_name] File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pymc\model_graph.py:77, in fast_eval(var) 76 def fast_eval(var): ---> 77 return function([], var, mode=_cheap_eval_mode)() File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pytensor\compile\function\types.py:1038, in Function.__call__(self, output_subset, *args, **kwargs) 1036 t0_fn = time.perf_counter() 1037 try: -> 1038 outputs = vm() if output_subset is None else vm(output_subset=output_subset) 1039 except Exception: 1040 self._restore_defaults() File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pytensor\link\utils.py:201, in streamline.<locals>.streamline_default_f() 199 old_s[0] = None 200 except Exception: --> 201 raise_with_op(fgraph, node, thunk) File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pytensor\link\utils.py:526, in raise_with_op(fgraph, node, thunk, exc_info, storage_map) 521 warnings.warn( 522 f"{exc_type} error does not allow us to add an extra error message" 523 ) 524 # Some exception need extra parameter in inputs. So forget the 525 # extra long error message in that case. --> 526 raise exc_value.with_traceback(exc_trace) File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pytensor\link\utils.py:197, in streamline.<locals>.streamline_default_f() 192 try: 193 # strict=False because we are in a hot loop 194 for thunk, node, old_storage in zip( 195 thunks, order, post_thunk_old_storage, strict=False 196 ): --> 197 thunk() 198 for old_s in old_storage: 199 old_s[0] = None File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pytensor\graph\op.py:545, in Op.make_py_thunk.<locals>.rval(p, i, o, n, cm) 537 @is_thunk_type 538 def rval( 539 p=p, (...) 543 cm=node_compute_map, 544 ): --> 545 r = p(n, [x[0] for x in i], o) 546 for entry in cm: 547 entry[0] = True File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pytensor\tensor\random\op.py:428, in RandomVariable.perform(self, node, inputs, outputs) 424 rng = deepcopy(rng) 426 outputs[0][0] = rng 427 outputs[1][0] = np.asarray( --> 428 self.rng_fn(rng, *args, None if size is None else tuple(size)), 429 dtype=self.dtype, 430 ) File ~\Downloads\mika-long.github.io\.venv\Lib\site-packages\pymc\distributions\custom.py:87, in CustomDistRV.rng_fn(cls, rng, *args) 85 args = list(args) 86 size = args.pop(-1) ---> 87 return cls._random_fn(*args, rng=rng, size=size) TypeError: genpoisson_rvs() got an unexpected keyword argument 'rng' Apply node that caused the error: CustomDist_y_rv{"(),()->()"}(RNG(<Generator(PCG64) at 0x1AA0248C580>), NoneConst{None}, theta, lam) Toposort index: 3 Inputs types: [RandomGeneratorType, <pytensor.tensor.type_other.NoneTypeT object at 0x000001AA77AC2900>, TensorType(float64, shape=()), TensorType(float64, shape=())] Inputs shapes: ['No shapes', 'No shapes', (), ()] Inputs strides: ['No strides', 'No strides', (), ()] Inputs values: [Generator(PCG64) at 0x1AA0248C580, None, array(0.85235818), array(-0.14067537)] Outputs clients: [[], [Shape(y)]] Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer): File "c:\Program Files\Positron\resources\app\extensions\positron-python\python_files\lib\ipykernel\py3\IPython\core\interactiveshell.py", line 3490, in run_ast_nodes if await self.run_code(code, result, async_=asy): File "c:\Program Files\Positron\resources\app\extensions\positron-python\python_files\lib\ipykernel\py3\IPython\core\interactiveshell.py", line 3550, in run_code exec(code_obj, self.user_global_ns, self.user_ns) File "C:\Users\slong\AppData\Local\Temp\ipykernel_36040\966756442.py", line 5, in <module> y = pm.CustomDist("y", theta, lam, logp=genpossion_logp, random=genpoisson_rvs) File "C:\Users\slong\Downloads\mika-long.github.io\.venv\Lib\site-packages\pymc\distributions\custom.py", line 733, in __new__ return _CustomDist( File "C:\Users\slong\Downloads\mika-long.github.io\.venv\Lib\site-packages\pymc\distributions\distribution.py", line 529, in __new__ rv_out = cls.dist(*args, **kwargs) File "C:\Users\slong\Downloads\mika-long.github.io\.venv\Lib\site-packages\pymc\distributions\custom.py", line 132, in dist return super().dist( File "C:\Users\slong\Downloads\mika-long.github.io\.venv\Lib\site-packages\pymc\distributions\distribution.py", line 598, in dist return cls.rv_op(*dist_params, size=create_size, **kwargs) File "C:\Users\slong\Downloads\mika-long.github.io\.venv\Lib\site-packages\pymc\distributions\custom.py", line 191, in rv_op return rv_op(*dist_params, **kwargs) HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
# def logp(value: TensorVariable, mu: TensorVariable) -> TensorVariable:
# return -(value - mu)**2
# def random(
# mu: np.ndarray | float,
# rng: Optional[np.random.Generator] = None,
# size : Optional[Tuple[int]]=None,
# ) -> np.ndarray | float :
# return rng.normal(loc=mu, scale=1, size=size)
# with pm.Model():
# mu = pm.Normal('mu', 0 , 1)
# pm.CustomDist(
# 'custom_dist',
# mu,
# logp=logp,
# random=random,
# observed=np.random.randn(100, 3),
# size=(100, 3),
# )
# prior = pm.sample_prior_predictive(10)Sampling: [custom_dist, mu]brms exampleWe can also use the example from brms to try to create a Beta-Binomial model. Note that this is also implemented natively in PyMC3, which allows us to double check and verify the result.
3 see https://www.pymc.io/projects/docs/en/latest/api/distributions/generated/pymc.BetaBinomial.html
The watermark extension is already loaded. To reload it, use:
%reload_ext watermark
Last updated: Fri Jan 09 2026
Python implementation: CPython
Python version : 3.13.5
IPython version : 8.18.1
pytensor: 2.35.1
xarray : 2025.12.0
arviz : 0.22.0
arviz_plots: 0.7.0
numpy : 2.3.5
platform : 1.0.8
pymc : 5.26.1
pytensor : 2.35.1
sys : 3.13.5 (main, Jun 12 2025, 12:42:35) [MSC v.1943 64 bit (AMD64)]
Watermark: 2.5.1