DOC: Some mistakes on pymc.CustomDist

[tegusi]

Issue with current documentation:

When I test https://www.pymc.io/projects/docs/en/latest/api/distributions/generated/pymc.CustomDist.html and its dist parameter with third code snippets, I met TypeError here(pymc version=5.0.2):

import pymc as pm
from pytensor.tensor import TensorVariable

def dist(
    lam: TensorVariable,
    shift: TensorVariable,
    size: TensorVariable,
) -> TensorVariable:
    return pm.Exponential.dist(lam, size=size) + shift

with pm.Model() as m:
    lam = pm.HalfNormal("lam")
    shift = -1
    pm.CustomDist(
        "custom_dist",
        lam,
        shift,
        dist=dist,
        observed=[-1, -1, 0],
    )

    prior = pm.sample_prior_predictive()
    posterior = pm.sample()

Error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[235], line 14
     12 lam = pm.HalfNormal("lam")
     13 shift = -1
---> 14 pm.CustomDist(
     15     "custom_dist",
     16     lam,
     17     shift,
     18     dist=dist,
     19     observed=[-1, -1, 0],
     20 )
     22 prior = pm.sample_prior_predictive()
     23 posterior = pm.sample()

File ~/miniconda/envs/canon/lib/python3.9/site-packages/pymc/distributions/distribution.py:984, in CustomDist.__new__(cls, name, random, logp, logcdf, moment, ndim_supp, ndims_params, dtype, *dist_params, **kwargs)
    972     return _CustomSymbolicDist(
    973         name,
    974         *dist_params,
   (...)
    981         **kwargs,
    982     )
    983 else:
--> 984     return _CustomDist(
    985         name,
    986         *dist_params,
    987         class_name=name,
    988         random=random,
    989         logp=logp,
    990         logcdf=logcdf,
    991         moment=moment,
    992         ndim_supp=ndim_supp,
    993         ndims_params=ndims_params,
    994         dtype=dtype,
    995         **kwargs,
    996     )
    997 return super().__new__(cls, name, *args, **kwargs)

File ~/miniconda/envs/canon/lib/python3.9/site-packages/pymc/distributions/distribution.py:308, in Distribution.__new__(cls, name, rng, dims, initval, observed, total_size, transform, *args, **kwargs)
    305     elif observed is not None:
    306         kwargs["shape"] = tuple(observed.shape)
--> 308 rv_out = cls.dist(*args, **kwargs)
    310 rv_out = model.register_rv(
    311     rv_out,
    312     name,
   (...)
    317     initval=initval,
    318 )
    320 # add in pretty-printing support

File ~/miniconda/envs/canon/lib/python3.9/site-packages/pymc/distributions/distribution.py:524, in _CustomDist.dist(cls, class_name, logp, logcdf, random, moment, ndim_supp, ndims_params, dtype, *dist_params, **kwargs)
    521 if random is None:
    522     random = default_not_implemented(class_name, "random")
--> 524 return super().dist(
    525     dist_params,
    526     class_name=class_name,
    527     logp=logp,
    528     logcdf=logcdf,
    529     random=random,
    530     moment=moment,
    531     ndim_supp=ndim_supp,
    532     ndims_params=ndims_params,
    533     dtype=dtype,
    534     **kwargs,
    535 )

File ~/miniconda/envs/canon/lib/python3.9/site-packages/pymc/distributions/distribution.py:385, in Distribution.dist(cls, dist_params, shape, **kwargs)
    383 ndim_supp = getattr(cls.rv_op, "ndim_supp", None)
    384 if ndim_supp is None:
--> 385     ndim_supp = cls.rv_op(*dist_params, **kwargs).owner.op.ndim_supp
    386 create_size = find_size(shape=shape, size=size, ndim_supp=ndim_supp)
    387 rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs)

File ~/miniconda/envs/canon/lib/python3.9/site-packages/pymc/distributions/distribution.py:579, in _CustomDist.rv_op(cls, class_name, logp, logcdf, random, moment, ndim_supp, ndims_params, dtype, *dist_params, **kwargs)
    576     return moment(rv, size, *dist_params)
    578 rv_op = rv_type()
--> 579 return rv_op(*dist_params, **kwargs)

File ~/miniconda/envs/canon/lib/python3.9/site-packages/pytensor/tensor/random/op.py:290, in RandomVariable.__call__(self, size, name, rng, dtype, *args, **kwargs)
    289 def __call__(self, *args, size=None, name=None, rng=None, dtype=None, **kwargs):
--> 290     res = super().__call__(rng, size, dtype, *args, **kwargs)
    292     if name is not None:
    293         res.name = name

File ~/miniconda/envs/canon/lib/python3.9/site-packages/pytensor/graph/op.py:296, in Op.__call__(self, *inputs, **kwargs)
    254 r"""Construct an `Apply` node using :meth:`Op.make_node` and return its outputs.
    255 
    256 This method is just a wrapper around :meth:`Op.make_node`.
   (...)
    293 
    294 """
    295 return_list = kwargs.pop("return_list", False)
--> 296 node = self.make_node(*inputs, **kwargs)
    298 if config.compute_test_value != "off":
    299     compute_test_value(node)

TypeError: make_node() got an unexpected keyword argument 'dist'

I am not quite familiar with dist function and its internal mechanism, please fix the function or code sample, thanks!

Idea or request for content:

No response

[ricardoV94]

You are looking at the developer documentation ("latest") for the next PyMC version (not yet released). The "stable" documentation shows the old API of 5.0.2: https://www.pymc.io/projects/docs/en/stable/api/distributions/generated/pymc.CustomDist.html where we were using the random kwarg for two different cases

[tegusi]

Thanks for your feedback, looking forward to the next version release.

You are looking at the developer documentation ("latest") for the next PyMC version (not yet released). The "stable" documentation shows the old API of 5.0.2: https://www.pymc.io/projects/docs/en/stable/api/distributions/generated/pymc.CustomDist.html where we were using the random kwarg for two different cases