Support logp derivation in DensityDist

[ricardoV94]

This PR allows users to return Aesara random graphs in the random method of DensityDist (now renamed to CustomDist), to be used in prior and posterior predictive sampling, as well as for automatic logp derivation using PyMC machinery.

For now, using this functionality will issue an experimental warning.

Closes #4530
Closes #4507

import pymc as pm

def shifted_exponential(lam, shift, size):
    return pm.Exponential.dist(lam, size=size) + shift

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

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

Here is an example that address this discourse help request: https://discourse.pymc.io/t/inference-with-aesara-pymc/10945

import pymc as pm

def pert(low, peak, high, lmbda, size):
    range = (high - low)
    s_alpha = 1 + lmbda * (peak - low) / range
    s_beta = 1 + lmbda * (high - peak) / range
    return pm.Beta.dist(s_alpha, s_beta, size=size) * range + low

with pm.Model() as m:
    low = pm.Normal("low", 0, 10)
    peak = pm.Normal("peak", 50, 10)
    high = pm.Normal("high", 100, 10)
    lmbda = 4
    pm.CustomDist("pert", low, peak, high, lmbda, random=pert, observed=[30.0, 35.0, 73.0])

m.point_logps()

Major / Breaking Changes

• ...

Bugfixes / New features

• CustomDist now accepts Aesara random graphs via the random kwarg. PyMC can sometimes infer the logp of these distributions without the need for a user to provide a logp function. This functionality is experimental!

Docs / Maintenance

• DensityDist is now an alias to CustomDist

[codecov]

Codecov Report

Merging #6361 (bcfff6d) into main (ddc6b65) will decrease coverage by 30.80%.
The diff coverage is 96.71%.

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #6361       +/-   ##
===========================================
- Coverage   94.71%   63.90%   -30.81%     
===========================================
  Files         132      131        -1     
  Lines       26695    26928      +233     
===========================================
- Hits        25284    17209     -8075     
- Misses       1411     9719     +8308     

Impacted Files	Coverage Δ
pymc/tests/distributions/test_logprob.py	0.00% <0.00%> (-100.00%)	⬇️
pymc/tests/sampling/test_forward.py	0.00% <0.00%> (-100.00%)	⬇️
pymc/pytensorf.py	84.59% <92.85%> (-10.04%)	⬇️
pymc/distributions/distribution.py	95.55% <98.23%> (+0.38%)	⬆️
pymc/tests/distributions/test_distribution.py	98.69% <99.24%> (+0.48%)	⬆️
pymc/exceptions.py	73.52% <100.00%> (-26.48%)	⬇️
pymc/model.py	89.56% <100.00%> (-0.21%)	⬇️
pymc/tests/sampling/test_parallel.py	81.48% <100.00%> (-1.49%)	⬇️
pymc/tests/smc/test_smc.py	0.00% <0.00%> (-100.00%)	⬇️
pymc/tests/backends/test_base.py	0.00% <0.00%> (-100.00%)	⬇️
... and 87 more

[zaxtax]

Do you think it might be good to have a stub DensityDist so users know they should migrate to CustomDist?

[ricardoV94]

Do you think it might be good to have a stub DensityDist so users know they should migrate to CustomDist?

It's there. Didn't even deprecate it because it's so commonly used.