benchmarking -- do not merge

Author: jessegrabowski

pytensor/tensor/slinalg.py

@@ -738,12 +738,20 @@ class LUSolve(Op):
     Solve a system of linear equations given the LU factorization of the matrix.
     """
 
-    __props__ = ("trans", "overwrite_b", "check_finite", "b_ndim")
+    __props__ = ("trans", "overwrite_b", "check_finite", "b_ndim", "expect_pivots")
 
-    def __init__(self, b_ndim, trans=False, overwrite_b=False, check_finite=True):
+    def __init__(
+        self,
+        b_ndim,
+        trans=False,
+        overwrite_b=False,
+        check_finite=True,
+        expect_pivots=False,
+    ):
         self.trans = trans
         self.overwrite_b = overwrite_b
         self.check_finite = check_finite
+        self.expect_pivots = expect_pivots
 
         assert b_ndim in (1, 2)
         self.b_ndim = b_ndim
@@ -789,6 +797,9 @@ def inplace_on_inputs(self, allowed_inplace_inputs: list[int]) -> "Op":
     def perform(self, node, inputs, outputs):
         LU, pivots, b = inputs
 
+        if not self.expect_pivots:
+            raise NotImplementedError
+
         outputs[0][0] = scipy_linalg.lu_solve(
             lu_and_piv=(LU, pivots),
             b=b,
@@ -807,8 +818,12 @@ def L_op(
         [x] = outputs
         [x_bar] = output_grads
 
-        p_inv = _pivot_to_permutation(pivots)
-        p = pt.argsort(p_inv)
+        if not self.expect_pivots:
+            p_inv = _pivot_to_permutation(pivots)
+            p = pt.argsort(p_inv)
+        else:
+            p = pivots
+
         P = ptb.identity_like(LU)[p]
 
         # We are solving PLUx = b

tests/tensor/test_slinalg.py

@@ -15,6 +15,7 @@
 from pytensor.tensor.slinalg import (
     Cholesky,
     CholeskySolve,
+    LUSolve,
     Solve,
     SolveBase,
     SolveTriangular,
@@ -703,7 +704,8 @@ def test_lu_factor(permutation_indices):
 
 @pytest.mark.parametrize("b_shape", [(5,), (5, 5)])
 @pytest.mark.parametrize("trans", [True, False])
-def test_lu_solve(b_shape: tuple[int], trans):
+@pytest.mark.parametrize("use_op", [True, False])
+def test_lu_solve(b_shape: tuple[int], trans, use_op):
     def T(x):
         if trans:
             return x.T
@@ -717,7 +719,13 @@ def T(x):
     b_val = rng.normal(size=b_shape).astype(config.floatX)
 
     LU_and_pivots = lu_factor(A)
-    x = lu_solve(LU_and_pivots, b, trans=trans)
+
+    if use_op:
+        x = LUSolve(b_ndim=len(b_shape), trans=trans, check_finite=False)(
+            LU_and_pivots, b
+        )
+    else:
+        x = lu_solve(LU_and_pivots, b, trans=trans)
 
     f = pytensor.function([A, b], x)
     x_pt = f(A_val.copy(), b_val.copy())
@@ -735,26 +743,6 @@ def T(x):
     )
     np.testing.assert_allclose(x_pt, x_sp)
 
-    # import jax
-    # import jax.scipy as jsp
-    #
-    # def jax_f(A, b):
-    #     LU_and_pivots = jsp.linalg.lu_factor(A)
-    #     x = jsp.linalg.lu_solve(LU_and_pivots, b, trans=trans)
-    #     return x.sum()
-
-    # jax_res = jax.value_and_grad(jax_f, [0, 1])(A_val, b_val)
-    # g = grad(x.sum(), [A, b])
-    # fg = pytensor.function([A, b], [x.sum(), *g])
-
-    # for a, b in zip(fg(A_val, b_val), [jax_res[0], *jax_res[1]]):
-    #     print(a - b)
-
-    # LU, pivots = pt.tensor('LU', shape=(5, 5)), pt.tensor('pivots', shape=(5,), dtype='int')
-    # x = lu_solve((LU, pivots), b)
-
-    # LU_val, pivots_val = scipy.linalg.lu_factor(A_val)
-
     utt.verify_grad(
         lambda A, b: lu_solve(lu_factor(A), b, trans=trans).sum(),
         pt=[A_val.copy(), b_val.copy()],
@@ -776,15 +764,6 @@ def test_fn(A, b):
         x = lu_solve(lu_and_pivots, b)
         return x.sum()
 
-    # A = pt.tensor("A", shape=(5, 5))
-    # b = pt.tensor("b", shape=b_shape)
-
-    # fg = pytensor.function([A, b], grad(test_fn(A, b), [A, b]))
-    # fg2 = pytensor.function([A, b], grad(pt.linalg.solve(A, b).sum(), [A, b]))
-
-    # print(fg(A_val, b_val))
-    # print(fg2(A_val, b_val))
-
     utt.verify_grad(test_fn, [A_val, b_val], 3, rng)
 
 
@@ -1065,3 +1044,33 @@ def test_block_diagonal_blockwise():
     B = np.random.normal(size=(1, batch_size, 4, 4)).astype(config.floatX)
     result = block_diag(A, B).eval()
     assert result.shape == (10, batch_size, 6, 6)
+
+
+def lu_solve_1(A, b):
+    lu, pivots = pt.linalg.lu_factor(A)
+    return pt.linalg.lu_solve((lu, pivots), b)
+
+
+def lu_solve_2(A, b, b_ndim=1, trans=0, check_finite=False):
+    lu, pivots = pt.linalg.lu_factor(A)
+    return LUSolve(b_ndim=1, trans=0, check_finite=False)(lu, pivots, b)
+
+
+@pytest.mark.parametrize(
+    "op", [lu_solve_1, lu_solve_2, pt.linalg.solve], ids=["lu_1", "lu_2", "solve"]
+)
+@pytest.mark.parametrize("n", [500])
+def test_solve_methods(op, n, benchmark):
+    A = pt.tensor("A", shape=(n, n))
+    b = pt.tensor("b", shape=(n,))
+
+    x = op(A, b)
+    gx = pt.grad(x.sum(), [A, b])
+    f = pytensor.function([A, b], [x, *gx])
+
+    A_val = np.random.normal(size=(n, n)).astype(config.floatX)
+    b_val = np.random.normal(size=(n,)).astype(config.floatX)
+
+    # Trigger compilation if we're a jit mode
+    f(A_val, b_val)
+    benchmark(f, A_val, b_val)