提交 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.graph.rewriting.basic import NodeRewriter
from pytensor.graph.rewriting.basic import NodeRewriter, in2out
from pytensor.graph.rewriting.db import EquilibriumDB, RewriteDatabase
from pytensor.tensor.rewriting.ofg import inline_ofg_expansion
lower_xtensor_db = EquilibriumDB(ignore_newtrees=False)
......@@ -14,6 +15,15 @@ optdb.register(
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(
node_rewriter: RewriteDatabase | NodeRewriter | str, *tags: str, **kwargs
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论