Merge pull request #506 from Lnaden/jax_stagger_jit

Wrap the preconditioned jit
choderalab · Jun 16, 2023 · cfe49fc · cfe49fc
2 parents 1aaa16d + 3598b01
commit cfe49fc
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/pymbar/mbar_solvers.py b/pymbar/mbar_solvers.py
@@ -1,5 +1,6 @@
 import logging
 import warnings
+from functools import wraps
 
 import numpy as np
 
@@ -135,6 +136,9 @@ def jit_or_pass_after_bitsize(jitable_fn):
         A function which can be jit'd
     """
 
+    @wraps(
+        jitable_fn
+    )  # Helper to ensure the decorated function still registers for docs and inspection
     def staggered_jit(*args, **kwargs):
         # This will only trigger if JAX is set
         if use_jit and not config.x64_enabled:

diff --git a/pymbar/tests/test_mbar_solvers.py b/pymbar/tests/test_mbar_solvers.py
@@ -64,7 +64,7 @@ def run_mbar_protocol(oscillator_bundle, protocol):
         "CG",
         "BFGS",
         "Newton-CG",
-        "TNC",
+        pytest.param("TNC", marks=pytest.mark.flaky(max_runs=2)),  # This one is flaky
         "trust-ncg",
         "trust-krylov",
         "trust-exact",