Add test for linear_solver kwarg

pymc-devs · Aug 17, 2021 · b13d83b · b13d83b
1 parent ec4a094
commit b13d83b
Showing 1 changed file with 23 additions and 0 deletions.
diff --git a/sunode/test_solve.py b/sunode/test_solve.py
@@ -152,3 +152,26 @@ def rhs(t, y, p):
 
     solver = AdjointSolver(problem)
     check_call_solve(solver, param_vals, "backward")
+
+
+def test_linear_solver_kwarg():
+    def rhs(t, y, p):
+        return {
+            'x': y.x,
+        }
+
+    states = {
+        'x': (),
+    }
+
+    params = {
+        'b': ()
+    }
+    param_vals = {
+        'b': 0.2
+    }
+    problem = SympyProblem(params, states, rhs, derivative_params=[])
+    linear_solver_opts = ["dense", "dense_finitediff", "spgmr_finitediff", "spgmr"]
+    for linear_solver in linear_solver_opts:
+        solver = Solver(problem, linear_solver=linear_solver)
+        check_call_solve(solver, param_vals, None)