Allow freezing of FunctionGraph for hashing

[jessegrabowski]

Closes #1606

LLM disclosure: this PR made heavy use of Claude in the planning and first cut stages, though I was heavily involved. Still, the code should be subject to extra scrutiny as a result.

The purpose of the PR is to refactor Ops with inner graphs to allow comparison. The linked issue has an exhaustive discussion of the factors at play. There was an attempt in the aesara days to attack this, but it was perhaps too aggressive: it cons-hashed all Apply nodes, which necessitated changes across the codebase. @ricardoV94 suggested a weakref dict approach for subgraphs. This is implemented at the Op level. The plan is for Ops that have inner graphs (Composite, ScalarLoop, Scan, OpFromGraph, etc) to have a _cache class attribute, and implement the op-specific logic for caching, pickling, unpickling, etc. It didn't look super generalizable to me at first blush, but we can argue about it maybe.

Changes to FunctionGraph:

• FunctionGraph now has a method freeze that returns a FrozenFunctionGraph.
• The FrozenFunctionGraph does cons-hashing of Apply nodes within its scope only
• It generates a hash based on its inner graph
• Two FrozenFunctionGraphs with the same inner graph with evaluate to equal, but their Apply nodes won't be references to the same objects (this is the "conservatism" of my approach)

Specific implementation details:

• The structural_hash of a FrozenFunctionGraph is built from a list of 3-tuples: (name, type, inputs), plus the outputs. For constants, inputs is replaced with the hash of the input data.
• Equality between FrozenFunctionGraphs is done by comparing hashes, then falling back to equal_computation if the hash misses.

A consequence of the cons-hashing in this approach is that the inner graph is de-duplicated when we call fg.freeze(). So a MergeOptimizer pass is no longer required. Usage is demonstrated on the Composite Op. If we like the approach I can move forward with refactoring other Ops, but I wanted to stop here and discuss the approach.

Code example:

import pytensor.tensor as pt
import pytensor

a, b, c, d = pt.dscalars('a', 'b', 'c', 'd')
eq1 = pt.sin(a) * b ** 2
eq2 = pt.sin(c) * d ** 2

with pytensor.config.change_flags(optimizer_verbose=True):
    f = pytensor.function([a, b, c, d], [eq1, eq2])

f.dprint()

Result:

Composite{(sin(*0-<float64>) * sqr(*1-<float64>))} [id A] 1
 ├─ a [id B]
 └─ b [id C]
Composite{(sin(*0-<float64>) * sqr(*1-<float64>))} [id D] 0
 ├─ c [id E]
 └─ d [id F]

Inner graphs:

Composite{(sin(*0-<float64>) * sqr(*1-<float64>))} [id A]
 ← mul [id G]
    ├─ sin [id H]
    │  └─ *0-<float64> [id I]
    └─ sqr [id J]
       └─ *1-<float64> [id K]

Composite{(sin(*0-<float64>) * sqr(*1-<float64>))} [id D]
 ← mul [id G]
    └─ ···

[ricardoV94]

Why did you not go all out?

If you already deduplicate and do internal hash-cons you are one step away from getting hashing for free across different FunctionGraphs. Just do the hash-cons globally. Then FrozenFunctionGrahp([x, y], [foo(x, y)] is equal to another functiongraph if and only if fgraph.outputs == other_fgraph.outputs. No need for recursive hashing or expensive equal_computations.

As it stands you are not doing much better sneaking a default MergeOptimizer at __init__ and adding a FunctionGraph class that has no replace mode.

And cheap hashing/ equality is not just a nice to have, it's really valuable to not slow down compilation. In some of my benchmarks on previous work, some graphs could spend inordinate time on equality checks.

Comments regardless of whether we go:

• Don't create FrozenFunctionGraph as a subclass of FrozenGraph, let's push the general principle, shared abstract classes, no-subclass of actually realized objects. Then you don't need check_frozen , the methods just don't exist for the frozen subclass.
• You could create a frozenApply that uses tuple for input/outputs instead of list. That will help ensuring the immutability because all our current rewrite machinery works on the idea of overriding entries in those lists. Accidentally trying to mutate a graph would 99% fail there.

[ricardoV94]

This is starting to look good, how are you feeling about it?

Notes:

• Add a FrozenFunctionGraph.unfreeze(), that yields a FunctionGraph?
• Really try to avoid the FrozenConstant stuff
• Ops with inner graph (at least the ones you touched now) should only have a FrozenFunctionGraph internally (not a mutable one as well). Maybe that's already the case.

We need some follow-up issues open:

• Optimizing OpFromGraph: There should be an explicit rewrite that creates a new OpFromGraph with its updated frozen graph, (so it is also reflected immediately in dprint). We should never do any further rewrites of the internal fgraph during compilation.
• Scan/Minimize/Root: Use the new FrozenFunctionGraph as well. This should immediately address #1601
• When compiling OpFromGraph in jitted contexts we should try to avoid recreating inner numba/jax functions when the same OFG is compiled multiple times in a function, this will likely speedup compilation. In the C-backend that already happens due to the caching of _fn. That's how we can deliver on the promised compilations speedups and it's specially relevant for a library like pytensor-ml that may want to chains hundreds of the same "LayerOp"s in sequence

[ricardoV94]

can't we have regular __props__ based equality/hashing now?

[ricardoV94]

@jessegrabowski this still stands. With proper fgraph equality, we could have these inner graph Ops behave like other Ops based on __props__ (simpler mental model for devs). __props__ = ("fgraph",) (and whatever else not in the fgraph that influences behavior) ?

[ricardoV94]

I left some comments as I checked the changes. I need to think/discuss a bit about the spec thing, and the desire to have a consistent hashing across runtimes. If you remove that the complexity of this PR drops quite a bit, but maybe this is also fine.

Can you confirm this was only needed for the C-backend, and that it would also work if whatever relies on that called something like __stable_hash__ instead of __hash__, that does the fingerprint / spec thing?

Besides that this PR look amazing, and it's a game changer to working with inner graph ops. We really need those to work well

[jessegrabowski]

I removed the spec stuff and simplified the PR down somewhat.

[ricardoV94]

Do we need the frozenapply to be hash-consed? Isn't it enough if the input/output variables are? Wondering if we can remove some extra code that way. The Apply doesn't do much anyway

[ricardoV94]

this is probably correct, but it's something that we have to think about. Some Ops have properties that affect their behavior in pytensor but not the computational meaning.

Those in theory should not be part of __props__ and affect equality. This would allow MergeOptimizer to merge the nodes with the same inputs but different OFG, which I guess is what we would want for the final compiled graph (e.g., it doesn't matter if node1 has a different gradient than node2 when you compile it at the end).

But on the other hand we don't want to confuse the two Ops, because when we "freeze -> unfreeze" for example, we don't want to lose those attributes.

We need to work on this in the future. There are different degrees of "equality" we need for different things. And maybe custom __eq__ even in the presence of __props__ is what we need, but I don't thing MergeOptimizer looks at __props__ specifically.

[ricardoV94]

This special behavior of shared variables is something I want to get rid of already for v3, but fine to test here as it's still a thing

[ricardoV94]

Wondering if in the rebuild_collect_shared that's used at the beginning of the function compilation, we will still merge the Op (if not the node ofc). Because that would be nice, only one compilation instead of 2. Just a curiosity, not something that needs to be tested here

[ricardoV94]

This is amazing. Small notes throughout, just want us to think on whether we also need the FrozenApply to be hash-consed. Maybe yes, I dunno why I'm drawing a distinction (I think I didn't like the extra pickling code needed).

[ricardoV94]

I think with this PR we'll stop seeing what was the actual optimized inner graph in the compiled function? Something we should follow up with, optimizing inner graph should be an explicit rewrite, not something that happens at make_node / dispatch time