Scan logprob fails when unvalued stochastic outputs are returned

Open
#6,909 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
42/100
Issue type
Bug
Clarity
Mostly clear
Activity status
Stale
Tech stack
python
Domain
backend

Research direction

Start with the Python reproduction in the issue, focusing on the two CustomDist definitions and their pytensor.scan outputs. Run it and compare logp1 with logp2; done means unvalued stochastic outputs no longer cause the two equivalent models to produce different log probabilities.

Written by the indexing model from the issue text.

Description

bug logprob
Description

Found by @lucianopaz

import numpy as np
import pymc as pm
import pytensor
import pytensor.tensor as pt

from pymc.pytensorf import collect_default_updates

steps = 4

def ar_dist1(rho, sigma, size):
    def ar_step(x_tm1, rho, sigma):
        eps_t = pm.Normal.dist(sigma=sigma)
        mu = x_tm1 * rho
        x = mu + eps_t
        return x, collect_default_updates([x])

    ar_innov, _ = pytensor.scan(
        fn=ar_step,
        outputs_info=[{"initial": pt.zeros(()), "taps": [-1]}],
        non_sequences=[rho, sigma],
        n_steps=steps,
        strict=True,
    )

    return ar_innov


def ar_dist2(rho, sigma, size):
    def ar_step(x_tm1, rho, sigma):
        eps_t = pm.Normal.dist(sigma=sigma)
        mu = x_tm1 * rho
        x = mu + eps_t
        return [x, eps_t], collect_default_updates([x])

    [ar_innov, _], _ = pytensor.scan(
        fn=ar_step,
        outputs_info=[{"initial": pt.zeros(()), "taps": [-1]}, None],
        non_sequences=[rho, sigma],
        n_steps=steps,
        strict=True,
    )

    return ar_innov


with pm.Model() as m:
    rho = 0.1
    sigma = 0.1
    observed = np.arange(steps)

    pm.CustomDist(
        "ar_dist1",
        rho,
        sigma,
        dist=ar_dist1,
        observed=observed,
    )

    pm.CustomDist(
        "ar_dist2",
        rho,
        sigma,
        dist=ar_dist2,
        observed=observed,
    )

logp1, logp2 = m.compile_logp(sum=False)({})
np.testing.assert_allclose(logp1, logp2)
"""
AssertionError: 
Not equal to tolerance rtol=1e-07, atol=0
Mismatched elements: 2 / 4 (50%)
Max absolute difference: 58.
Max relative difference: 0.12928641
 x: array([   1.383647,  -48.616353, -179.116353, -390.616353])
 y: array([   1.383647,  -48.616353, -198.616353, -448.616353])
"""
Dominant language
Python
Stars
9.8k
Forks
2.3k
Avg merge
21h 39m
Merged PRs (30d)
5

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from pymc-devs/pymc

All issues in pymc-devs/pymc

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.