Extend logprob inference for scans with carried auxiliary states

[ricardoV94]

PyMC can now infer graphs like MA(2)

import numpy as np
import pytensor

import pymc as pm

init_eps = pm.DiracDelta.dist(np.zeros(2, dtype="float64"))
rho = [0.3, 0.9]
sigma = 1
n_steps = 100

def ma2_step(eps_tm2, eps_tm1, rho, sigma):
    mu = eps_tm1 * rho[0] + eps_tm2 * rho[1]
    y = pm.Normal.dist(mu, sigma)
    eps = y - mu
    return eps, y

[_, ma2], _ = pytensor.scan(
    fn=ma2_step,
    outputs_info=[{"initial": init_eps, "taps": range(-2, 0)}, None],
    non_sequences=[rho, sigma],
    n_steps=n_steps,
    name="ma2",
)

ma2_test = ma2.eval()
pm.logp(ma2, ma2_test).sum().eval()

CC @jessegrabowski

[codecov]

Codecov Report

Merging #6582 (eddcc8f) into main (c709abf) will decrease coverage by 14.28%.
The diff coverage is 12.50%.

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #6582       +/-   ##
===========================================
- Coverage   92.02%   77.75%   -14.28%     
===========================================
  Files          92       92               
  Lines       15563    15592       +29     
===========================================
- Hits        14322    12123     -2199     
- Misses       1241     3469     +2228     

Impacted Files	Coverage Δ
pymc/logprob/scan.py	17.67% <7.40%> (-79.25%)	⬇️
pymc/testing.py	82.19% <23.07%> (-9.65%)	⬇️

... and 53 files with indirect coverage changes

[larryshamalama]

Ouf, Scans a bit beyond my area of comfort. I just have minor comments for this one

[twiecki]

Need to add tags so this shows up in release notes.

[ricardoV94]

Need to add tags so this shows up in release notes.

Already had. Without tags it still shows up in the release notes, but in the catch-all "maintenance" bucket. The only case it doesn't show up is if it has the label "no releasenotes"