Test numba solve gradients

Author: jessegrabowski

pytensor/link/numba/dispatch/_LAPACK.py

@@ -279,7 +279,7 @@ def numba_xgetrf(cls, dtype):
     def numba_xgetrs(cls, dtype):
         """
         Solve a system of linear equations A @ X = B or A.T @ X = B with a general N-by-N matrix A using the LU
-        factorization computed by numba_getrf.
+        factorization computed by GETRF.
 
         Called by scipy.linalg.lu_solve
         """
@@ -302,8 +302,8 @@ def numba_xgetrs(cls, dtype):
     @classmethod
     def numba_xsysv(cls, dtype):
         """
-        Solve a system of linear equations A @ X = B with a symmetric matrix A using the factorization computed by
-        sytrf (LDL or UDU).
+        Solve a system of linear equations A @ X = B with a symmetric matrix A using the diagonal pivoting method,
+        factorizing A into LDL^T or UDU^T form, depending on the value of UPLO
 
         Called by scipy.linalg.solve when assume_a == "sym"
         """
@@ -327,10 +327,8 @@ def numba_xsysv(cls, dtype):
     @classmethod
     def numba_xsycon(cls, dtype):
         """
-        Estimates the reciprocal of the condition number of a symmetric matrix A using the factorization computed by
-        sytrf (LDL or UDU).
-
-        Called by scipy.linalg.solve when assume_a == "sym"
+        Estimate the reciprocal of the condition number of a symmetric matrix A using the UDU or LDL factorization
+        computed by xSYTRF.
         """
         lapack_ptr, float_pointer = _get_lapack_ptr_and_ptr_type(dtype, "sycon")
 

tests/link/numba/test_slinalg.py

@@ -1,4 +1,5 @@
 import re
+from functools import partial
 
 import numpy as np
 import pytest
@@ -162,7 +163,7 @@ def test_numba_Cholesky_raise_on(on_error):
 @pytest.mark.parametrize("trans", [True, False], ids=["trans=True", "trans=False"])
 def test_numba_Cholesky_grad(lower, trans):
     rng = np.random.default_rng(utt.fetch_seed())
-    L = rng.random.normal(size=(5, 5)).astype(floatX)
+    L = rng.normal(size=(5, 5)).astype(floatX)
     X = L @ L.T
 
     utt.verify_grad(pt.linalg.cholesky, [X])
@@ -358,9 +359,11 @@ def test_solve(b_func, b_size, assume_a, transposed):
 
     X_np = f(A_val, b_val)
     op = f.maker.fgraph.outputs[0].owner.op
+
     # overwrite_b is preferred when both inputs can be destroyed
     assert op.destroy_map == {0: [1]}
 
+    # Test that the result is numerically correct
     np.testing.assert_allclose(
         transpose_func(A_val_copy, transposed) @ X_np, b_val_copy, atol=ATOL, rtol=RTOL
     )
@@ -370,6 +373,15 @@ def test_solve(b_func, b_size, assume_a, transposed):
     assert (b_val == b_val_copy).all() == (op.destroy_map.get(0, None) != [1])
 
     # Test gradients
+    solve = partial(
+        pt.linalg.solve,
+        lower=False,
+        assume_a=assume_a,
+        transposed=transposed,
+        b_ndim=len(b_size),
+    )
+
+    utt.verify_grad(solve, [A_val_copy, b_val_copy], mode="NUMBA")
 
 
 @pytest.mark.parametrize(