Skip to content

Commit 4c40efa

Browse files
committed
Stabilize rewrites shouldn't target Dot22
1 parent 2aa4aad commit 4c40efa

File tree

1 file changed

+3
-4
lines changed

1 file changed

+3
-4
lines changed

pytensor/tensor/rewriting/linalg.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
diag,
2424
diagonal,
2525
)
26-
from pytensor.tensor.blas import Dot22
2726
from pytensor.tensor.blockwise import Blockwise
2827
from pytensor.tensor.elemwise import DimShuffle, Elemwise
2928
from pytensor.tensor.math import Dot, Prod, _matmul, log, outer, prod
@@ -103,12 +102,12 @@ def transinv_to_invtrans(fgraph, node):
103102

104103

105104
@register_stabilize
106-
@node_rewriter([Dot, Dot22])
105+
@node_rewriter([Dot])
107106
def inv_as_solve(fgraph, node):
108107
"""
109108
This utilizes a boolean `symmetric` tag on the matrices.
110109
"""
111-
if isinstance(node.op, Dot | Dot22):
110+
if isinstance(node.op, Dot):
112111
l, r = node.inputs
113112
if (
114113
l.owner
@@ -280,7 +279,7 @@ def cholesky_ldotlt(fgraph, node):
280279
A.owner is not None
281280
and (
282281
(
283-
isinstance(A.owner.op, Dot | Dot22)
282+
isinstance(A.owner.op, Dot)
284283
# This rewrite only applies to matrix Dot
285284
and A.owner.inputs[0].type.ndim == 2
286285
)

0 commit comments

Comments
 (0)