Added rewrite for diag of kronecker product

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -12,14 +12,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,
@@ -701,3 +703,20 @@ def rewrite_inv_diag_to_diag_reciprocal(fgraph, node):
         non_eye_input = pt.shape_padaxis(non_eye_diag, -2)
 
     return [eye_input / non_eye_input]
+
+
+@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

@@ -662,3 +662,26 @@ def test_inv_diag_from_diag(inv_op):
         atol=ATOL,
         rtol=RTOL,
     )
+
+
+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,
+    )