Skip to content

Commit 7096b3c

Browse files
committed
Stabilize rewrites shouldn't target Dot22
1 parent f0fb328 commit 7096b3c

File tree

1 file changed

+3
-11
lines changed

1 file changed

+3
-11
lines changed

pytensor/tensor/rewriting/linalg.py

Lines changed: 3 additions & 11 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
@@ -278,14 +277,7 @@ def cholesky_ldotlt(fgraph, node):
278277
A = node.inputs[0]
279278
if not (
280279
A.owner is not None
281-
and (
282-
(
283-
isinstance(A.owner.op, Dot | Dot22)
284-
# This rewrite only applies to matrix Dot
285-
and A.owner.inputs[0].type.ndim == 2
286-
)
287-
or (A.owner.op == _matmul)
288-
)
280+
and ((isinstance(A.owner.op, Dot)) or (A.owner.op == _matmul))
289281
):
290282
return
291283

0 commit comments

Comments
 (0)