pymc-devs/pymc

CheckParameterValues not replaced in inner graphs

Open

#6,697 opened on Apr 28, 2023

View on GitHub
 (3 comments) (0 reactions) (0 assignees)Python (7,926 stars) (1,902 forks)batch import
help wantedlogprob

Description

We include these rewrites when compiling logp graphs, but this won't replace inner logp graphs (e.g, inside Scan or not-inlined OpFromGraph)

https://github.com/pymc-devs/pymc/blob/371472dcf1c44f99a9ad4e7317606404b2a47a0e/pymc/pytensorf.py#L1142-L1144

The example in #6696 suffers from this. Sampling is quite faster when the rewrite is manually added

Contributor guide

Tech stack
python
Domain
backendperformance
Issue type
bug
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
4
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
half day
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
stale
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
clear
Prerequisites
Understanding of PyMC compilation pipelineFamiliarity with graph rewrites in PyTensor
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
30
Research direction
The issue is that the CheckParameterValues rewrite is not applied to inner graphs (e.g., inside Scan or not inlined OpFromGraph). The relevant code is at pymc/pytensorf.py lines 1142-1144. The fix would involve ensuring this rewrite is recursively applied to all subgraphs. Start by examining how rewrites are applied in the compilation process, and look at the example in issue #6696 to reproduce the performance difference. Consider modifying the rewrite application logic to traverse inner graphs.