Wrap a diffeqzoo.backend.select to avoid unnecessary ci-crashes

pnkraemer · Jun 13, 2024 · adffcb1 · adffcb1
1 parent 8e406a9
commit adffcb1
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/docs/benchmarks/taylor_node/run_taylor_node.py b/docs/benchmarks/taylor_node/run_taylor_node.py
@@ -165,7 +165,10 @@ def adaptive_benchmark(fun, *, timeit_fun: Callable, max_time) -> dict:
 
 if __name__ == "__main__":
     set_jax_config()
-    backend.select("jax")
+
+    if not backend.has_been_selected:
+        backend.select("jax")
+
     algorithms = {
         r"Forward-mode": odejet_via_jvp(),
         r"Taylor-mode (scan)": taylor_mode_scan(),