Added rewrite for diag of kronecker product

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -11,14 +11,16 @@
 from pytensor.scalar.basic import Mul
 from pytensor.tensor.basic import (
     AllocDiag,
+    ExtractDiag,
     Eye,
     TensorVariable,
+    diag,
     diagonal,
 )
 from pytensor.tensor.blas import Dot22
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
-from pytensor.tensor.math import Dot, Prod, _matrix_matrix_matmul, log, prod
+from pytensor.tensor.math import Dot, Prod, _matrix_matrix_matmul, log, outer, prod
 from pytensor.tensor.nlinalg import (
     SVD,
     KroneckerProduct,
@@ -611,3 +613,20 @@ def rewrite_inv_inv(fgraph, node):
     ):
         return None
     return [potential_inner_inv.inputs[0]]
+
+
+@register_canonicalize
+@register_stabilize
+@node_rewriter([ExtractDiag])
+def rewrite_diag_kronecker(fgraph, node):
+    # Check for inner kron operation
+    potential_kron = node.inputs[0].owner
+    if not (potential_kron and isinstance(potential_kron.op, KroneckerProduct)):
+        return None
+
+    # Find the matrices
+    a, b = potential_kron.inputs
+    diag_a, diag_b = diag(a), diag(b)
+    outer_prod_as_vector = outer(diag_a, diag_b).flatten()
+
+    return [outer_prod_as_vector]

tests/tensor/rewriting/test_linalg.py

@@ -568,3 +568,26 @@ def get_pt_function(x, op_name):
     op2 = get_pt_function(op1, inv_op_2)
     rewritten_out = rewrite_graph(op2)
     assert rewritten_out == x
+
+
+def test_diag_kronecker_rewrite():
+    a, b = pt.dmatrices("a", "b")
+    kron_prod = pt.linalg.kron(a, b)
+    diag_kron_prod = pt.diag(kron_prod)
+    f_rewritten = function([a, b], diag_kron_prod, mode="FAST_RUN")
+
+    # Rewrite Test
+    nodes = f_rewritten.maker.fgraph.apply_nodes
+    assert not any(isinstance(node.op, KroneckerProduct) for node in nodes)
+
+    # Value Test
+    a_test, b_test = np.random.rand(2, 20, 20)
+    kron_prod_test = np.kron(a_test, b_test)
+    diag_kron_prod_test = np.diag(kron_prod_test)
+    rewritten_val = f_rewritten(a_test, b_test)
+    assert_allclose(
+        diag_kron_prod_test,
+        rewritten_val,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )