Add return type hint

Author: jessegrabowski

pytensor/tensor/rewriting/linalg.py

@@ -1,7 +1,9 @@
 import logging
+from collections.abc import Callable
 from typing import cast
 
-from pytensor.graph import FunctionGraph, Apply
+from pytensor import Variable
+from pytensor.graph import Apply, FunctionGraph
 from pytensor.graph.rewriting.basic import copy_stack_trace, node_rewriter
 from pytensor.tensor.basic import TensorVariable, diagonal
 from pytensor.tensor.blas import Dot22
@@ -320,7 +322,9 @@ def local_log_prod_sqr(fgraph, node):
 
 @register_specialize
 @node_rewriter([Blockwise])
-def local_lift_through_linalg(fgraph: FunctionGraph, node: Apply):
+def local_lift_through_linalg(
+    fgraph: FunctionGraph, node: Apply
+) -> list[Variable] | None:
     """
     Rewrite compositions of linear algebra operations by lifting expensive operations (Cholesky, Inverse) through Ops
     that join matrices (KroneckerProduct, BlockDiagonal).
@@ -340,7 +344,7 @@ def local_lift_through_linalg(fgraph: FunctionGraph, node: Apply):
 
     Returns
     -------
-    res: list of Variable, optional
+    list of Variable, optional
         List of optimized variables, or None if no optimization was performed
     """
 
@@ -357,15 +361,15 @@ def local_lift_through_linalg(fgraph: FunctionGraph, node: Apply):
             input_matrices = y.owner.inputs
 
             if isinstance(outer_op.core_op, MatrixInverse):
-                outer_f = inv
+                outer_f = cast(Callable, inv)
             elif isinstance(outer_op.core_op, Cholesky):
-                outer_f = cholesky
+                outer_f = cast(Callable, cholesky)
             elif isinstance(outer_op.core_op, MatrixPinv):
-                outer_f = pinv
+                outer_f = cast(Callable, pinv)
             else:
                 raise NotImplementedError
 
-            inner_matrices = [outer_f(m) for m in input_matrices]
+            inner_matrices = [cast(TensorVariable, outer_f(m)) for m in input_matrices]
 
             if isinstance(y.owner.op, KroneckerProduct):
                 return [kron(*inner_matrices)]