Improve docstring

Author: jessegrabowski

pytensor/tensor/nlinalg.py

@@ -1017,8 +1017,6 @@ class KroneckerProduct(OpFromGraph):
     Wrapper Op for Kronecker graphs
     """
 
-    ...
-
 
 def kron(a, b):
     """Kronecker product.

pytensor/tensor/rewriting/linalg.py

@@ -1,6 +1,7 @@
 import logging
 from typing import cast
 
+from pytensor.graph import FunctionGraph, Apply
 from pytensor.graph.rewriting.basic import copy_stack_trace, node_rewriter
 from pytensor.tensor.basic import TensorVariable, diagonal, swapaxes
 from pytensor.tensor.blas import Dot22
@@ -322,37 +323,32 @@ def local_log_prod_sqr(fgraph, node):
         # 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):
+def local_lift_through_linalg(fgraph: FunctionGraph, node: Apply):
     """
-    Rewrite a graph like Inv(BlockDiag([A, B, C])) to BlockDiag([Inv(A), Inv(B), Inv(C)])
+    Rewrite compositions of linear algebra operations by lifting expensive operations (Cholesky, Inverse) through Ops
+    that join matrices (KroneckerProduct, BlockDiagonal).
+
+    This rewrite takes advantage of commutation between certain linear algebra operations to do several smaller matrix
+    operations on component matrices instead of one large one. For example, when taking the inverse of Kronecker
+    product, we can take the inverse of each component matrix and then take the Kronecker product of the inverses. This
+    reduces the cost of the inverse from O((n*m)^3) to O(n^3 + m^3) where n and m are the dimensions of the component
+    matrices.
 
     Parameters
     ----------
-    fgraph
-    node
+    fgraph: FunctionGraph
+        Function graph being optimized
+    node: Apply
+        Node of the function graph to be optimized
 
     Returns
     -------
-
+    res: list of Variable, optional
+        List of optimized variables, or None if no optimization was performed
     """
+
     # TODO: Simplify this if we end up Blockwising KroneckerProduct
     if isinstance(node.op.core_op, MatrixInverse | Cholesky | MatrixPinv):
         y = node.inputs[0]

tests/tensor/rewriting/test_linalg.py

@@ -386,6 +386,4 @@ def test_local_lift_through_linalg(constructor, f_op, f, g_op, g):
     ]
     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), atol=1e-8)