backend: promote _under_jax_trace -> _namespaces.under_jax_trace (shared trace predicate)

The jit/grad/vmap tracer check that gates the eager-loop-vs-lax.fori_loop
choice is general (bracket expansion, bisection, any future rolled loop), not
a root-finder detail, so it moves next to is_backend_array/device_of in
galpy.backend._namespaces and both call sites in optimize.py plus the
duplicate inline check in _jax/optimize._bisect_root now share it.

Also add a jit+grad-through-Staeckel test: Staeckel's turning points call the
module-level bisect_root DIRECTLY (not via brentq), so its under_jax_trace(a,b)
-> lax.fori_loop branch was uncovered (the Adiabatic/Spherical/Vertical grad
tests route through brentq/iterate_bracket instead). The new test traces
actionsFreqsAngles under jax.jit (asserts the jaxpr is rolled) and jax.grad
(finite dJr/dR), closing the codecov/patch gap on optimize.py:219-222.

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: jobovy

galpy/backend/_jax/optimize.py

@@ -49,9 +49,10 @@ def _bisect_root(f, a, b, xp, *, xtol, maxiter):
     """
     import jax
 
+    from .._namespaces import under_jax_trace
     from ..optimize import bisect_root, bisect_step, n_bisect_steps
 
-    if not any(isinstance(x, jax.core.Tracer) for x in (a, b)):
+    if not under_jax_trace(a, b):  # entered directly with a concrete bracket
         return bisect_root(f, a, b, xp, xtol=xtol, maxiter=maxiter)
     lo = xp.asarray(a) * 1.0
     hi = xp.asarray(b) * 1.0

galpy/backend/_namespaces.py

@@ -47,6 +47,31 @@ def is_backend_array(x):
     return False
 
 
+def under_jax_trace(*xs):
+    """True iff jax is imported AND one of ``xs`` is a jax tracer (jit/grad/vmap).
+
+    The predicate that gates the eager-loop-vs-``lax.fori_loop`` choice wherever
+    galpy rolls a fixed-schedule loop (bracket expansion, bisection, ...): the
+    eager Python loop stays byte-identical and ~9x faster outside a trace, while
+    under a jax trace the same body is rolled into a ``fori_loop`` so its ``n``
+    embedded copies of the physics closure do not unroll into the user's jaxpr.
+
+    Cheap on numpy/torch and on plain (untraced) jax arrays: if ``jax`` is not
+    even imported we short-circuit to ``False`` (via ``sys.modules``, so the
+    numpy/torch eager paths never import jax). This is deliberately gated on
+    ``sys.modules`` rather than the ``_JAX_LOADED`` install flag so a jax-
+    installed-but-unused run (pure numpy/torch) keeps the eager hot path from
+    importing jax at all.
+    """
+    import sys
+
+    if "jax" not in sys.modules:
+        return False
+    import jax
+
+    return any(isinstance(x, jax.core.Tracer) for x in xs)
+
+
 def _is_floating_dtype(dtype):
     """True for real floating-point dtypes of any backend.
 

galpy/backend/optimize.py

@@ -46,7 +46,7 @@
 #   first-order sensitivity; a Hessian through x* would require a true
 #   custom_jvp differentiating the implicit relation a second time.
 ###############################################################################
-from ._namespaces import is_backend_array
+from ._namespaces import is_backend_array, under_jax_trace
 from ._resolver import get_namespace
 
 # Default bracketing tolerance (matches scipy.optimize.brentq's xtol default, so
@@ -156,23 +156,6 @@ def bisect_step(lo, hi, slo, f, xp):
     return xp.where(same, mid, lo), xp.where(same, hi, mid)
 
 
-def _under_jax_trace(*xs):
-    """True iff jax is imported AND one of ``xs`` is a jax tracer (jit/grad/vmap).
-
-    Cheap on numpy/torch and on plain jax arrays: if ``jax`` is not even imported
-    we short-circuit to ``False`` (so the eager Python loops stay byte-identical
-    and we never import jax on the numpy/torch path). Under a jax trace the n
-    copies of the physics closure would otherwise unroll into the user's jaxpr.
-    """
-    import sys
-
-    if "jax" not in sys.modules:
-        return False
-    import jax
-
-    return any(isinstance(x, jax.core.Tracer) for x in xs)
-
-
 def iterate_bracket(step, x0, n):
     """Run ``x = step(x)`` ``n`` times -- a fixed-schedule, branch-free bracket
     expansion/contraction where ``step`` is a single ``xp.where`` update over the
@@ -183,7 +166,7 @@ def iterate_bracket(step, x0, n):
     ``step`` so the ``n`` embedded copies of ``f`` do NOT unroll into the jaxpr
     (mirrors the bisection rolling; eager stays ~9x faster than ``fori_loop``).
     """
-    if _under_jax_trace(x0):
+    if under_jax_trace(x0):
         import jax
 
         return jax.lax.fori_loop(0, n, lambda _, x: step(x), x0)
@@ -216,7 +199,7 @@ def bisect_root(f, a, b, xp, *, xtol, maxiter):
     # DIRECT callers jit-fast too -- notably the Staeckel umin/umax/vmin turning
     # points, which call bisect_root straight (not via brentq). Eager keeps the
     # Python loop (bit-identical, ~9x faster than fori_loop).
-    if _under_jax_trace(a, b):
+    if under_jax_trace(a, b):
         from ._jax.optimize import _bisect_root
 
         return _bisect_root(f, a, b, xp, xtol=xtol, maxiter=maxiter)

tests/test_backend_actionAngle.py

@@ -841,6 +841,30 @@ def test_staeckel_actions_vs_c(backend):
     numpy.testing.assert_allclose(_np(jz_b), numpy.asarray(jz_c), rtol=1e-8, atol=1e-9)
 
 
+@pytest.mark.skipif("jax" not in BACKENDS, reason="jax not installed")
+def test_staeckel_jit_grad_rolls_direct_bisection():
+    # Staeckel's turning points call the module-level ``bisect_root`` DIRECTLY
+    # (not via ``brentq``). Under a jax trace (jit/grad) the bracket endpoints are
+    # tracers, so ``bisect_root`` dispatches to the rolled ``lax.fori_loop`` kernel
+    # -- the ~100-step bisection does NOT unroll ~100 copies of the Staeckel
+    # integrand into the user's jaxpr. Covers the ``under_jax_trace(a, b)`` branch
+    # of ``galpy.backend.optimize.bisect_root`` reached only by these direct
+    # callers (the brentq-based AA methods route through a different kernel).
+    aA = actionAngleStaeckel(pot=MWPotential2014, delta=0.45, c=False)
+    R = jnp.asarray(_STK[0])
+    rest = tuple(jnp.asarray(v) for v in _STK[1:])
+    jr_e, _, jz_e = aA(R, *rest)  # eager (Python-loop) reference
+    jr_j, _, jz_j = jax.jit(lambda r: aA(r, *rest))(R)  # traced (fori_loop) value
+    numpy.testing.assert_allclose(_np(jr_j), _np(jr_e), rtol=1e-8, atol=1e-10)
+    numpy.testing.assert_allclose(_np(jz_j), _np(jz_e), rtol=1e-8, atol=1e-10)
+    # the jaxpr is ROLLED: a loop primitive, not ~100 unrolled bisection steps.
+    txt = str(jax.make_jaxpr(lambda r: aA(r, *rest)[0])(R))
+    assert ("while" in txt) or ("scan" in txt)
+    # grad flows through the direct-bisection turning points: finite dJr/dR.
+    g = jax.grad(lambda r: jnp.sum(aA(r, *rest)[0]))(R)
+    assert numpy.all(numpy.isfinite(_np(g)))
+
+
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_staeckel_unbound_backend_no_raise(backend):
     # An unbound orbit raises UnboundError on the numpy path (eager), but must NOT