Default rewrites in inner graph

[aerubanov]

This PR will add implementation of default rewrites for nodes with inner fgraph and related with issue #6697. For now I just added new test cases

---

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

[codecov]

Codecov Report

Merging #6996 (b8fde47) into main (547bcb4) will decrease coverage by 4.25%.
Report is 11 commits behind head on main.
The diff coverage is 94.11%.

❗ Current head b8fde47 differs from pull request most recent head 0c80a68. Consider uploading reports for the commit 0c80a68 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6996      +/-   ##
==========================================
- Coverage   92.17%   87.92%   -4.25%     
==========================================
  Files         101      101              
  Lines       16849    16891      +42     
==========================================
- Hits        15530    14851     -679     
- Misses       1319     2040     +721     

Files	Coverage Δ
pymc/logprob/utils.py	96.19% <94.11%> (-1.31%)	⬇️

... and 8 files with indirect coverage changes

[ricardoV94]

The tests are great!

[aerubanov]

Logic for removing CheckParameterValues from Scan is very similar to one we recently added in #6873 for moment calculation in Scan, so may be I will move it to some helper function later

[ricardoV94]

There is some refactoring going on in #6976 that will impact the location of these changes.

One thing we need to think about is nested Scans/OpFromGraph, I think this approach will only work for the first level.

PS: Agree with refactoring the functionality if it's so similar

[aerubanov]

I added recursive check for nested Scans. For now it is only remove CheckParameterValues nodes from fgraph, but I will add replacement by ninf next step

[aerubanov]

@ricardoV94 I added logic for replacement CheckParameterValues by ninf in Scan and rebased my branch. Now I am going to add support for OpFromGraph and want to extract common logic for rewrites in inner fgraphs - probably some helper function. I think we can add some general solution for inner fgraphs rewriting.

[aerubanov]

I added support for replacements in OpFromGraph

[aerubanov]

I move some common logic into InnerGraphRewrite class

[aerubanov]

@ricardoV94 Could you please take a look when you will have some time?

[ricardoV94]

This is great. Functionality wise I think it's complete. I just have some suggestions to make it more similar to other PyTensor rewriters, and some test organization. Feel free to push back on my suggestions.

[ricardoV94]

This may be overdoing it, but what do you think of implementing this as something like WalkingGraphRewriter: https://github.com/ricardoV94/pytensor/blob/9cf2d181f07dc99bbd2e7c9e2b4a3e1b0aeff034/pytensor/graph/rewriting/basic.py#L1998

A WalkingNestedGraphRewriter which applies the same node_rewriter to both the outer graph and inner graphs. The idea is you would pass the previous rewrite which doesn't distinguish between the core case or a Scan OpFromGraph. It would be the WalkingNestedGraphRewriter that would apply that logic regardless of which NodeRewriter it's given?

[ricardoV94]

If not, this should at least inherit from NodeRewriter, which already takes care of enforcing a self.transform is implemented, so you don't have to.

[aerubanov]

Hmm, sounds great. I've been thinking about something like that, not sure how to best implement it.

[ricardoV94]

This class should be implemented in pytensorf. The specific rewrites that use it can be implemented here though

[ricardoV94]

Instead of messing with models and compile_pymc, let's just use pytensor.function directly and pass the rewrites we want.

pytensor.function(..., mode=get_mode().including("rewrite_name"))

These tests should be in logprob/test_utils.py

[aerubanov]

I will change and move tests

[aerubanov]

I will re-implement this with WalkingNestedGraphRewriter after merging pymc-devs/pytensor#556