Added dot kron rewrite

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -42,6 +42,7 @@
     register_specialize,
     register_stabilize,
 )
+from pytensor.tensor.shape import Reshape
 from pytensor.tensor.slinalg import (
     BlockDiagonal,
     Cholesky,
@@ -989,3 +990,23 @@ def jax_bilinaer_lyapunov_to_direct(fgraph: FunctionGraph, node: Apply):
     "jax",
     position=0.9,  # Run before canonicalization
 )
+
+
+@register_canonicalize
+@register_stabilize
+@node_rewriter([Reshape])
+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
+
+    c = node.inputs[1]
+    [a, b] = potential_kron.inputs
+
+    m, n = a.type.shape
+    p, q = b.type.shape
+    out_clever = (b @ c.reshape(shape=(n, q)).T @ a.T).ravel()
+    return [out_clever]

tests/tensor/rewriting/test_linalg.py

@@ -906,3 +906,29 @@ def test_rewrite_cholesky_diag_to_sqrt_diag_not_applied():
     f_rewritten = function([x], z_cholesky, mode="FAST_RUN")
     nodes = f_rewritten.maker.fgraph.apply_nodes
     assert any(isinstance(node.op, Cholesky) for node in nodes)
+
+
+def test_dot_kron_rewrite():
+    m, n, p, q = 3, 4, 6, 7
+    a = pt.matrix("a", shape=(m, n))
+    b = pt.matrix("b", shape=(p, q))
+    c = pt.matrix("c", shape=(n * q, 1))
+    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
+
+    # NUMERIC VALUE TEST
+    a_test = np.random.rand(m, n)
+    b_test = np.random.rand(p, q)
+    c_test = np.random.rand(n * q, 1)
+    out_direct_val = np.kron(a_test, b_test) @ c_test
+    out_clever_val = f_direct_rewritten(a_test, b_test, c_test)
+    assert_allclose(
+        out_direct_val,
+        out_clever_val,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )