tests: adds calculate_sensitivities check for JaxSolver

pybamm-team · Sep 26, 2024 · 83b5710 · 83b5710
1 parent 48fff9a
commit 83b5710
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tests/unit/test_solvers/test_jax_solver.py b/tests/unit/test_solvers/test_jax_solver.py
@@ -44,6 +44,11 @@ def test_model_solver(self):
 
             np.testing.assert_array_equal(second_solution.y, solution.y)
 
+            # Test passing `calculate_sensitivities`
+            solution_sens = solver.solve(model, t_eval, calculate_sensitivities=True)
+            np.testing.assert_array_equal(solution_sens.y, solution.y)
+            assert len(solution_sens.sensitivities) == 0
+
     def test_semi_explicit_model(self):
         # Create model
         model = pybamm.BaseModel()