Support Scans in CustomDist

[ricardoV94]

Here is how one can now write an AR with StudentT innovations:

import numpy as np
import pymc as pm
import pytensor.tensor as pt
from pymc.pytensorf import collect_default_updates
from pytensor.scan import scan

def ar_t(nu, sigma, rho, size):
    
    def step(xtm1, nu, sigma, rho):
        x = xtm1 * rho + pm.StudentT.dist(nu=nu, sigma=sigma)
        return x, collect_default_updates([x])
    
    xs, _ = scan(
        fn=step,
        outputs_info=[pt.zeros(())],
        non_sequences=[nu, sigma, rho],
        n_steps=size[0],
    )
    
    return xs

with pm.Model() as m:
    nu = 4
    sigma = pm.HalfNormal("sigma")
    rho = pm.Uniform("rho")
    steps=100
    
    pm.CustomDist("ar_t", nu, sigma, rho, dist=ar_t, observed=np.random.randn(steps))
    
    prior = pm.sample_prior_predictive()
    posterior = pm.sample()    

The part that was not supported was that we would not be able to collect the update for the scan (used in prior/posterior predictive), and the user had no way of specifying it manually either.

---

📚 Documentation preview 📚: https://pymc--6696.org.readthedocs.build/en/6696/

[michaelosthege]

When reviewing I slowed down at find_default_update inside collect_default_update. This entire local function can be extracted which would make it a lot easier to review & test.

Overall it looks good, but I'd recommend to add some comments to make it comprehensible without knowing the context of this PR.

[ricardoV94]

This entire local function can be extracted which would make it a lot easier to review & test.

It's the inner recursive function, I don't think it makes sense to be outside of the initial caller function.

[codecov]

Codecov Report

Merging #6696 (6fa3a31) into main (c57769c) will decrease coverage by 10.99%.
The diff coverage is 100.00%.

❗ Current head 6fa3a31 differs from pull request most recent head 5532cb2. Consider uploading reports for the commit 5532cb2 to get more accurate results

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #6696       +/-   ##
===========================================
- Coverage   92.01%   81.02%   -10.99%     
===========================================
  Files          95       95               
  Lines       16180    16176        -4     
===========================================
- Hits        14888    13107     -1781     
- Misses       1292     3069     +1777     

Impacted Files	Coverage Δ
pymc/distributions/distribution.py	95.23% <100.00%> (-1.41%)	⬇️
pymc/pytensorf.py	90.67% <100.00%> (-1.93%)	⬇️

... and 44 files with indirect coverage changes

[jessegrabowski]

I think I understand what's going on here. The key change is just to add a check for Scan Ops and, if so, reach into the outer_out_from_outer_inp mapping and grab the rng object, raising an error if it's not found.

I also found the function inside recursive function a bit unusual. This SO thread claims there is some overhead associated with doing things this way, so I'm wondering if that would matter in a worst-case recursion (very big graph? multiple scans?).

Otherwise looks good to me. I would second a bit of documentation in the functions -- a general complaint I have about the Theano codebase is a lack of useful docstrings.

[ricardoV94]

I think I understand what's going on here. The key change is just to add a check for Scan Ops and, if so, reach into the outer_out_from_outer_inp mapping and grab the rng object, raising an error if it's not found.

Yup

I also found the function inside recursive function a bit unusual. This SO thread claims there is some overhead associated with doing things this way, so I'm wondering if that would matter in a worst-case recursion (very big graph? multiple scans?).

I think that's only meaningful when you are calling the outer function many times, in this case the outer function is called once and then the inner function is called multiple times. Also this is not performance critical code at all since it's done once before compiling a pytensor function. The compilation is orders of magnitude more expensive than this one time graph recursion. And this compilation is itself order of magnitude faster than the presumed heavy use of the final compiled function.

The comments in that thread seem to agree with this. One user also mentions the uncluttering of the module namespace which is the main motivation I had.

Otherwise looks good to me. I would second a bit of documentation in the functions -- a general complaint I have about the Theano codebase is a lack of useful docstrings.

Sure

[ricardoV94]

Updated docstrings and added one code example