changing tracked node to dot

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -42,7 +42,6 @@
     register_specialize,
     register_stabilize,
 )
-from pytensor.tensor.shape import Reshape
 from pytensor.tensor.slinalg import (
     BlockDiagonal,
     Cholesky,
@@ -994,11 +993,8 @@ def jax_bilinaer_lyapunov_to_direct(fgraph: FunctionGraph, node: Apply):
 
 @register_canonicalize
 @register_stabilize
-@node_rewriter([Reshape])
+@node_rewriter([Dot])
 def rewrite_dot_kron(fgraph, node):
-    if not (isinstance(node.op, Blockwise) and isinstance(node.op.core_op, Dot)):
-        return False
-
     potential_kron = node.inputs[0].owner
     if not (isinstance(potential_kron.op, KroneckerProduct)):
         return False

tests/tensor/rewriting/test_linalg.py

@@ -14,7 +14,7 @@
 from pytensor.tensor import swapaxes
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle
-from pytensor.tensor.math import _allclose, dot, matmul
+from pytensor.tensor.math import Dot, _allclose, dot, matmul
 from pytensor.tensor.nlinalg import (
     SVD,
     Det,
@@ -916,9 +916,10 @@ def test_dot_kron_rewrite():
     out_direct = pt.linalg.kron(a, b) @ c
 
     # REWRITE TEST
-    f_direct_rewritten = function([a, b, c], out_direct)
-    # nodes = f_direct_rewritten.maker.fgraph.apply_nodes
-    # Add assertion test here
+    f_direct_rewritten = function([a, b, c], out_direct, mode="FAST_RUN")
+    nodes = f_direct_rewritten.maker.fgraph.apply_nodes
+    print(nodes)
+    assert not any(isinstance(node.op.core_op, Dot) for node in nodes)
 
     # NUMERIC VALUE TEST
     a_test = np.random.rand(m, n)