File tree Expand file tree Collapse file tree 1 file changed +11
-1
lines changed
pytensor/xtensor/rewriting Expand file tree Collapse file tree 1 file changed +11
-1
lines changed Original file line number Diff line number Diff line change 1
1
from pytensor .compile import optdb
2
- from pytensor .graph .rewriting .basic import NodeRewriter
2
+ from pytensor .graph .rewriting .basic import NodeRewriter , in2out
3
3
from pytensor .graph .rewriting .db import EquilibriumDB , RewriteDatabase
4
+ from pytensor .tensor .rewriting .ofg import inline_ofg_expansion
4
5
5
6
6
7
lower_xtensor_db = EquilibriumDB (ignore_newtrees = False )
14
15
position = 0.1 ,
15
16
)
16
17
18
+ # Register OFG inline again after lowering xtensor
19
+ optdb .register (
20
+ "inline_ofg_expansion_xtensor" ,
21
+ in2out (inline_ofg_expansion ),
22
+ "fast_run" ,
23
+ "fast_compile" ,
24
+ position = 0.11 ,
25
+ )
26
+
17
27
18
28
def register_lower_xtensor (
19
29
node_rewriter : RewriteDatabase | NodeRewriter | str , * tags : str , ** kwargs
You can’t perform that action at this time.
0 commit comments