提交 13e52d87 authored 作者: Ricardo Vieira's avatar Ricardo Vieira 提交者: Ricardo Vieira

Add a second inline_ofg_expansion in xtensor for Ops that wrap OpFromGraph once lowered

上级 5fff8fff
from pytensor.compile import optdb from pytensor.compile import optdb
from pytensor.graph.rewriting.basic import NodeRewriter from pytensor.graph.rewriting.basic import NodeRewriter, in2out
from pytensor.graph.rewriting.db import EquilibriumDB, RewriteDatabase from pytensor.graph.rewriting.db import EquilibriumDB, RewriteDatabase
from pytensor.tensor.rewriting.ofg import inline_ofg_expansion
lower_xtensor_db = EquilibriumDB(ignore_newtrees=False) lower_xtensor_db = EquilibriumDB(ignore_newtrees=False)
...@@ -14,6 +15,15 @@ optdb.register( ...@@ -14,6 +15,15 @@ optdb.register(
position=0.1, position=0.1,
) )
# Register OFG inline again after lowering xtensor
optdb.register(
"inline_ofg_expansion_xtensor",
in2out(inline_ofg_expansion),
"fast_run",
"fast_compile",
position=0.11,
)
def register_lower_xtensor( def register_lower_xtensor(
node_rewriter: RewriteDatabase | NodeRewriter | str, *tags: str, **kwargs node_rewriter: RewriteDatabase | NodeRewriter | str, *tags: str, **kwargs
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论