Add rewrite to lift linear algebra through certain linalg ops

Author: jessegrabowski

pytensor/compile/builders.py

@@ -6,7 +6,7 @@
 from typing import Optional, cast
 
 import pytensor.tensor as pt
-from pytensor import function
+from pytensor.compile import function
 from pytensor.compile.function.pfunc import rebuild_collect_shared
 from pytensor.compile.mode import optdb
 from pytensor.compile.sharedvalue import SharedVariable

pytensor/tensor/nlinalg.py

@@ -600,7 +600,7 @@ def svd(a, full_matrices: bool = True, compute_uv: bool = True):
 
     Returns
     -------
-    U, V,  D : matrices
+    U, V, D : matrices
 
     """
     return SVD(full_matrices, compute_uv)(a)

pytensor/tensor/rewriting/linalg.py

@@ -7,17 +7,21 @@
 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, det
+from pytensor.tensor.nlinalg import MatrixInverse, MatrixPinv, det, inv, pinv
 from pytensor.tensor.rewriting.basic import (
     register_canonicalize,
     register_specialize,
     register_stabilize,
 )
 from pytensor.tensor.slinalg import (
+    BlockDiagonal,
     Cholesky,
+    KroneckerProduct,
     Solve,
     SolveBase,
+    block_diag,
     cholesky,
+    kron,
     solve,
     solve_triangular,
 )
@@ -310,3 +314,65 @@ def local_log_prod_sqr(fgraph, node):
 
         # TODO: have a reduction like prod and sum that simply
         # returns the sign of the prod multiplication.
+
+
+def local_inv_kron_to_kron_inv(fgraph, node):
+    # check if we have a kron
+    # check if parent node is an inv
+    # if yes, replace with kron(inv, inv)
+
+    pass
+
+
+def local_chol_kron_to_kron_chol(fgraph, node):
+    # check if we have a kron
+    # check if parent node is a cholesky
+    # if yes, replace with kron(cholesky, cholesky)
+
+    pass
+
+
+@register_specialize
+@node_rewriter([Blockwise])
+def local_lift_through_linalg(fgraph, node):
+    """
+    Rewrite a graph like Inv(BlockDiag([A, B, C])) to BlockDiag([Inv(A), Inv(B), Inv(C)])
+
+    Parameters
+    ----------
+    fgraph
+    node
+
+    Returns
+    -------
+
+    """
+    # TODO: Simplify this if we end up Blockwising KroneckerProduct
+    if isinstance(node.op.core_op, (MatrixInverse, Cholesky, MatrixPinv)):
+        y = node.inputs[0]
+        outer_op = node.op
+
+        if y.owner and (
+            isinstance(y.owner.op, Blockwise)
+            and isinstance(y.owner.op.core_op, BlockDiagonal)
+            or isinstance(y.owner.op, KroneckerProduct)
+        ):
+            input_matrices = y.owner.inputs
+
+            if isinstance(outer_op.core_op, MatrixInverse):
+                outer_f = inv
+            elif isinstance(outer_op.core_op, Cholesky):
+                outer_f = cholesky
+            elif isinstance(outer_op.core_op, MatrixPinv):
+                outer_f = pinv
+            else:
+                raise NotImplementedError
+
+            inner_matrices = [outer_f(m) for m in input_matrices]
+
+            if isinstance(y.owner.op, KroneckerProduct):
+                return [kron(*inner_matrices)]
+            elif isinstance(y.owner.op.core_op, BlockDiagonal):
+                return [block_diag(*inner_matrices)]
+            else:
+                raise NotImplementedError

pytensor/tensor/slinalg.py

@@ -9,6 +9,7 @@
 
 import pytensor
 import pytensor.tensor as pt
+from pytensor.compile.builders import OpFromGraph
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
 from pytensor.tensor import as_tensor_variable
@@ -559,6 +560,14 @@ def eigvalsh(a, b, lower=True):
     return Eigvalsh(lower)(a, b)
 
 
+class KroneckerProduct(OpFromGraph):
+    """
+    Wrapper Op for Kronecker graphs
+    """
+
+    ...
+
+
 def kron(a, b):
     """Kronecker product.
 
@@ -578,10 +587,13 @@ def kron(a, b):
     numpy.kron(a, b) != scipy.linalg.kron(a, b)!
     They don't have the same shape and order when
     a.ndim != b.ndim != 2.
-
     """
+
+    # TODO: Revisit this implementation?
+
     a = as_tensor_variable(a)
     b = as_tensor_variable(b)
+
     if a.ndim + b.ndim <= 2:
         raise TypeError(
             "kron: inputs dimensions must sum to 3 or more. "
@@ -598,7 +610,8 @@ def kron(a, b):
             (o.shape[0] * o.shape[2], o.shape[1] * o.shape[3])
             + tuple(o.shape[i] for i in range(4, o.ndim))
         )
-    return o
+
+    return KroneckerProduct(inputs=[a, b], outputs=[o])(a, b)
 
 
 class Expm(Op):

tests/tensor/rewriting/test_linalg.py

@@ -14,10 +14,12 @@
 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, matrix_inverse
+from pytensor.tensor.nlinalg import Det, MatrixInverse, MatrixPinv, matrix_inverse
 from pytensor.tensor.rewriting.linalg import inv_as_solve
 from pytensor.tensor.slinalg import (
+    BlockDiagonal,
     Cholesky,
+    KroneckerProduct,
     Solve,
     SolveBase,
     SolveTriangular,
@@ -333,3 +335,52 @@ def test_invalid_batched_a(self):
             ref_fn(test_a, test_b),
             rtol=1e-7 if config.floatX == "float64" else 1e-5,
         )
+
+
+@pytest.mark.parametrize(
+    "constructor", [pt.dmatrix, pt.tensor3], ids=["not_batched", "batched"]
+)
+@pytest.mark.parametrize(
+    "f_op, f",
+    [
+        (MatrixInverse, pt.linalg.inv),
+        (Cholesky, pt.linalg.cholesky),
+        (MatrixPinv, pt.linalg.pinv),
+    ],
+    ids=["inv", "cholesky", "pinv"],
+)
+@pytest.mark.parametrize(
+    "g_op, g",
+    [(BlockDiagonal, pt.linalg.block_diag), (KroneckerProduct, pt.linalg.kron)],
+    ids=["block_diag", "kron"],
+)
+def test_local_lift_through_linalg(constructor, f_op, f, g_op, g):
+    A, B = list(map(constructor, "ab"))
+    X = f(g(A, B))
+
+    f1 = pytensor.function(
+        [A, B], X, mode=get_default_mode().including("local_lift_through_linalg")
+    )
+    f2 = pytensor.function(
+        [A, B], X, mode=get_default_mode().excluding("local_lift_through_linalg")
+    )
+
+    all_apply_nodes = f1.maker.fgraph.apply_nodes
+    f_ops = [
+        x for x in all_apply_nodes if isinstance(getattr(x.op, "core_op", x.op), f_op)
+    ]
+    g_ops = [
+        x for x in all_apply_nodes if isinstance(getattr(x.op, "core_op", x.op), g_op)
+    ]
+
+    assert len(f_ops) == 2
+    assert len(g_ops) == 1
+
+    test_vals = [
+        np.random.normal(size=(3,) * A.ndim).astype(config.floatX) for _ in range(2)
+    ]
+    test_vals = [x @ np.swapaxes(x, -1, -2) for x in test_vals]
+
+    f2(*test_vals)
+
+    np.testing.assert_allclose(f1(*test_vals), f2(*test_vals))