Added rewrite for diag of kronecker product

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -22,7 +22,7 @@
 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,
@@ -818,3 +818,20 @@ def rewrite_slogdet_blockdiag(fgraph, node):
     )
 
     return [prod(sign_sub_matrices), sum(logdet_sub_matrices)]
+
+
+@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

@@ -751,3 +751,26 @@ def test_slogdet_blockdiag_rewrite():
         atol=1e-3 if config.floatX == "float32" else 1e-8,
         rtol=1e-3 if config.floatX == "float32" else 1e-8,
     )
+
+
+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,
+    )