Import kron and KroneckerProduct from nlinalg

Author: jessegrabowski

pytensor/tensor/nlinalg.py

@@ -7,6 +7,7 @@
 from numpy.core.numeric import normalize_axis_tuple  # type: ignore
 
 from pytensor import scalar as ps
+from pytensor.compile.builders import OpFromGraph
 from pytensor.gradient import DisconnectedType
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
@@ -1011,6 +1012,14 @@ def tensorsolve(a, b, axes=None):
     return TensorSolve(axes)(a, b)
 
 
+class KroneckerProduct(OpFromGraph):
+    """
+    Wrapper Op for Kronecker graphs
+    """
+
+    ...
+
+
 def kron(a, b):
     """Kronecker product.
 
@@ -1042,7 +1051,8 @@ def kron(a, b):
     out_shape = tuple(a.shape * b.shape)
     output_out_of_shape = a_reshaped * b_reshaped
     output_reshaped = output_out_of_shape.reshape(out_shape)
-    return output_reshaped
+
+    return KroneckerProduct(inputs=[a, b], outputs=[output_reshaped])(a, b)
 
 
 __all__ = [

pytensor/tensor/rewriting/linalg.py

@@ -7,7 +7,15 @@
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle
 from pytensor.tensor.math import Dot, Prod, _matrix_matrix_matmul, log, prod
-from pytensor.tensor.nlinalg import MatrixInverse, MatrixPinv, det, inv, pinv
+from pytensor.tensor.nlinalg import (
+    KroneckerProduct,
+    MatrixInverse,
+    MatrixPinv,
+    det,
+    inv,
+    kron,
+    pinv,
+)
 from pytensor.tensor.rewriting.basic import (
     register_canonicalize,
     register_specialize,
@@ -16,12 +24,10 @@
 from pytensor.tensor.slinalg import (
     BlockDiagonal,
     Cholesky,
-    KroneckerProduct,
     Solve,
     SolveBase,
     block_diag,
     cholesky,
-    kron,
     solve,
     solve_triangular,
 )
@@ -348,7 +354,7 @@ def local_lift_through_linalg(fgraph, node):
 
     """
     # TODO: Simplify this if we end up Blockwising KroneckerProduct
-    if isinstance(node.op.core_op, (MatrixInverse, Cholesky, MatrixPinv)):
+    if isinstance(node.op.core_op, MatrixInverse | Cholesky | MatrixPinv):
         y = node.inputs[0]
         outer_op = node.op
 

tests/tensor/rewriting/test_linalg.py

@@ -14,12 +14,17 @@
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle
 from pytensor.tensor.math import _allclose, dot, matmul
-from pytensor.tensor.nlinalg import Det, MatrixInverse, MatrixPinv, matrix_inverse
+from pytensor.tensor.nlinalg import (
+    Det,
+    KroneckerProduct,
+    MatrixInverse,
+    MatrixPinv,
+    matrix_inverse,
+)
 from pytensor.tensor.rewriting.linalg import inv_as_solve
 from pytensor.tensor.slinalg import (
     BlockDiagonal,
     Cholesky,
-    KroneckerProduct,
     Solve,
     SolveBase,
     SolveTriangular,