Skip to content

Do not recompile identical OFG in same graph#2209

Merged
ricardoV94 merged 5 commits into
pymc-devs:mainfrom
ricardoV94:ofg_single_dprint_and_compile
Jun 15, 2026
Merged

Do not recompile identical OFG in same graph#2209
ricardoV94 merged 5 commits into
pymc-devs:mainfrom
ricardoV94:ofg_single_dprint_and_compile

Conversation

@ricardoV94

@ricardoV94 ricardoV94 commented Jun 10, 2026

Copy link
Copy Markdown
Member

Draft: as it includes commits from #2056

Repeated OpFromGraph make a lot of sense for deep machine learning graphs. This PR makes more use of this functionality.

In the JIT backends: don't redefine / rewrite the same fgraph.

Numba compile is usually slower than our rewrites, and our caching mechanism already "deduplicates" the functions numba sees. So this is only a %speedup never much absolute. (Big speedup with cache=False, but we are not optimizing for that if it was our goal)


Also when dprint don't show identical copies:

import pytensor
import pytensor.tensor as pt

x = pt.vector("x")
relu_op = pytensor.OpFromGraph([x], [pt.maximum(x, 0)], inline=False, name="Relu")

x = pt.vector("x")
a = relu_op(x)
b = relu_op(a)
c = relu_op(b)
out = pt.exp(c)

out.dprint()

Before:

Exp [id A]
 └─ Relu{inline=False} [id B]
    └─ Relu{inline=False} [id C]
       └─ Relu{inline=False} [id D]
          └─ x [id E]

Inner graphs:

Relu{inline=False} [id B]
 ← Maximum [id F]
    ├─ i0 [id G]
    └─ ExpandDims{axis=0} [id H]
       └─ 0 [id I]

Relu{inline=False} [id C]
 ← Maximum [id F]
    └─ ···

Relu{inline=False} [id D]
 ← Maximum [id F]
    └─ ···

After:

Exp [id A]
 └─ Relu{inline=False} [id B]
    └─ Relu{inline=False} [id C]
       └─ Relu{inline=False} [id D]
          └─ x [id E]

Inner graphs:

Relu{inline=False} [id B, C, D]
 ← Maximum [id F]
    ├─ i0 [id G]
    └─ ExpandDims{axis=0} [id H]
       └─ 0 [id I]

@ricardoV94 ricardoV94 force-pushed the ofg_single_dprint_and_compile branch from 5bd4f4a to d464497 Compare June 10, 2026 10:00
@ricardoV94 ricardoV94 marked this pull request as draft June 10, 2026 10:03
@ricardoV94 ricardoV94 force-pushed the ofg_single_dprint_and_compile branch 2 times, most recently from c80bda8 to 9c40d08 Compare June 10, 2026 11:40
@ricardoV94 ricardoV94 force-pushed the ofg_single_dprint_and_compile branch from 9c40d08 to 1f83063 Compare June 15, 2026 11:21
… graphs

clone_v_get_shared_updates recursed once per graph node, overflowing Python's
call stack on graphs deeper than the recursion limit. Rewrite it as an
explicit-stack depth-first traversal with identical semantics.
@ricardoV94 ricardoV94 force-pushed the ofg_single_dprint_and_compile branch from 1f83063 to f9f8def Compare June 15, 2026 12:07
@ricardoV94 ricardoV94 marked this pull request as ready for review June 15, 2026 14:45
@ricardoV94 ricardoV94 merged commit 17239d0 into pymc-devs:main Jun 15, 2026
66 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants