Fix MergeOptimizer's dead guard against merging destroyed variables

[velochy]

The bug

MergeOptimizer.apply has a guard meant to skip merging identical nodes whose outputs
are destroyed by inplace clients:

if any(
    i in flatten(c.op.destroy_map.values())
    for c, i in clients
    if c.op.destroy_map
):
    continue

pytensor.utils.flatten only recurses into tuple | list | set. dict.values() is a
dict_values view, so flatten returns it wrapped in a single-element list and the
membership test can never succeed:

>>> flatten({0: [0]}.values())
[dict_values([[0]])]
>>> 0 in flatten({0: [0]}.values())
False

Under Python 2, dict.values() returned a list and the guard worked. It has been dead
code since the Python 3 port.

Why it matters

With the guard dead, every merge of two destroyed variables runs the full
replace_all_validate cycle — checkpoint, DestroyHandler validation failure
(InconsistencyError: Multiple destroyers), revert — and MergeFeature re-schedules the
pair after every subsequent graph change, so the same doomed pair is retried across the
whole rewrite phase.

Graphs with many identical destroyed nodes hit this combinatorially. The canonical shape
is the gradient of sum_i params_i[idx_i] (any additive model over gathered parameter
blocks): each term contributes a zeros-Alloc gradient buffer, all identical, and after
the inplace pass each is destroyed by its own AdvancedIncSubtensor{inplace}. Failed
merge attempts then grow ~n³:

n terms	failed merge attempts	pytensor.function (mode=NUMBA, main)	fixed
10	210	0.6 s	—
20	1,520	4.4 s	1.3 s (0 failures)
40	11,440	294.5 s	2.1 s (0 failures)

The blowup is backend-independent (FAST_RUN produces identical failure counts); numba
only makes the wall-clock more visible. A real 142-variable PyMC model spent 40+ minutes
and tens of GB in this storm without finishing.

Repro script

import numpy as np, pytensor, pytensor.tensor as pt

n, N, L, K = 40, 5232, 8, 20
rng = np.random.default_rng(0)
x = pt.dvector("x")
mu = pt.zeros((N, K))
for i in range(n):
    block = x[i*L*K:(i+1)*L*K].reshape((L, K))
    mu = mu + block[pt.constant(rng.integers(0, L, N))]
loss = pt.special.log_softmax(mu, axis=-1)[
    pt.arange(N), pt.constant(rng.integers(0, K, N))].sum()
g = pytensor.grad(loss, x)
f = pytensor.function([x], [loss, g], mode="NUMBA")   # 294s before, 2.1s after

The fix

Skip a scheduled pair when both variables have a direct destroyer:

if _has_direct_destroyer(fgraph, pairs[0][0]) and _has_direct_destroyer(
    fgraph, pairs[0][1]
):
    continue

Merging two destroyed variables always fails validation ("multiple destroyers"), so
these pairs are pure waste. Everything else is left to DestroyHandler validation,
which is the only thing that can decide it: a single-destroyer merge may or may not be
feasible (reordering can be possible, or can be proven cyclic — there is no cheap local
rule), and destroyers reachable only through a view chain are rare enough that paying
the occasional validate+revert beats recomputing the transitive droot map per pair. The
check is O(direct clients) and touches no DestroyHandler state.

The now-unused flatten import is removed.

Numerical impact

None by construction: the guard only skips replacements that DestroyHandler validation
would have rejected and reverted. Verified bit-identical logp and grad (max rel. diff
~1e-16, float noise) on a 60-RV hierarchical PyMC model compiled before/after.

Test

TestMergeOptimizer::test_no_merge_attempt_on_destroyed_variables: two identical op2
nodes each destroyed by a destroy_map={0: [0]} client; asserts nb_fail == 0 from the
MergeOptimizer profile and that the twins stay unmerged. Red on main (nb_fail == 2),
green with the fix.

[ricardoV94]

Under Python 2, dict.values() returned a list and the guard worked. It has been dead
code since the Python 3 port.

🤦 that's borderline hilarious. Also means we had 0 check coverage.

A single-destroyer merge is valid (the destroyer migrates to the surviving variable) and stays allowed

I don't think this is true at all? Maybe I'm just misreading, trivial case:

x  = pt.vector()
e1 = pt.exp(x) 
e2 = pt.exp(x)
out1 = e[0].set_inplace(0)  # destroy e1
out2 = e2[::-1]  # consumes e2

function([x], [out1 out2])

e1 and e2 can't be merged, our out2 would be corrupted.

out2 is a view, so it makes it obviously wrong, whereas if it's not a view it may still be possible to reorder out2 before out1, but this may create a cycle and be proven invalid. The point is, there's no local trivial rule to guarantee a merge is valid that's assymptotically cheaper than what the DestroyHandler / rejection mechanism ultimately runs. We can just try to get less false positives.

TLDR: The code is right (it doesn't reject as it may still be possible to fuse), that sentence itself is wrong (it's not guaranteed that the merging will be feasible).

also catches variables destroyed through a view chain, which a direct-clients destroy_map scan cannot see;

I worry about the performance of fgraph.has_destroyers. Is that cached? It's definitely better to not try mergings that will be rejected, but OTOH most times it's not this case. The choice here is not the flatten options vs the more comprehensive check, but the best narrow-scope we can do vs the comprehensive check. Just need to look at this with skepticism to see if we still conclude the same.

[ricardoV94]

Awesome find btw!

[velochy]

You're right on both counts — the sentence was wrong (there is no cheap local rule for the single-destroyer case; validation is the decider) and I've rewritten the description accordingly.

On the perf concern: fgraph.destroyers() is only cached until the next graph change (stale_droot), so per-pair transitive checks interleaved with successful merges would recompute the droot map repeatedly. Switched to a direct-clients scan (_has_direct_destroyer, O(clients), no DestroyHandler state touched), applied only in the both-destroyed case — which is the always-rejected one, and the one the pathological graphs hit: their destroyer is a direct client, so the storm benchmark is unchanged (n=40: 294.5s → 2.3s, 0 failed merges). View-chain-only destroyers and single-destroyer pairs fall through to validation as before.

[ricardoV94]

inline this function in the apply method (can still be a function), not in the loop of it

[velochy]

Moved into apply, above the loop.

[ricardoV94]

Can you also add test cases where one is destroyed and the other is not (which CAN be achieved with reordering) and the one where it cannot (the non destroyer needs both the twin and the op that destroys it)?

[ricardoV94]

unless they already exist elsewhere

[velochy]

Added both: test_merge_single_destroyer_feasible (merge goes through, reader ordered before the destroyer) and test_merge_single_destroyer_infeasible (reader also depends on the destroyer's result — validation rejects, twins stay distinct; the attempt is made and counted in nb_fail, documenting that single-destroyer pairs are validation's call, not the guard's). Didn't find existing coverage of the MergeOptimizer/DestroyHandler interplay — tests/graph/test_destroyhandler.py tests validation and orderings but never runs the merger.