Gradient assembly JIT-compiles an XLA kernel (jax.grad) for a log scale's scalar derivative, making pybnf[jax] mandatory for most gradient fits

[wshlavacek]

assemble_* differentiates a log-scaled parameter's sampling-space transform with jax.grad, JIT-compiling an XLA kernel to obtain a scalar derivative that is one line of algebra. Every gradient fit whose parameters are declared loguniform_var pays it, and every such fit hard-requires the pybnf[jax] extra.

Where

pybnf/gradient/assembly.py:

def _d_theta_d_u(free_param, u):
    """``d theta/d u`` for one log-scaled parameter via autodiff of ``inverse_jax``."""
    jax = _require_jax()
    return float(jax.grad(free_param.from_sampling_space_jax)(float(u)))

from_sampling_space_jax for a log scale is u -> 10**u (or exp(u)). Its derivative is closed form:

d(10^u)/du  = 10^u * ln(10)
d(exp(u))/du = exp(u)

Both are already available as the transform's own value times a constant, so the derivative needs no autodiff at all.

Why it matters

It is on the hot path. _d_theta_d_u is called per free parameter to build the native→sampling Jacobian, which is rebuilt as the fit proceeds. The XLA trace/compile shows up in the logs of an ordinary fit:

DEBUG jax._src.compiler:108: PERSISTENT COMPILATION CACHE MISS for 'jit__power'
      with key 'jit__power-20ef157d99ca469c1a1c9b0c803d08e8fa2cce1e450bfcebfa33b4d79bc44812'
DEBUG jax._src.compiler:723: Not writing persistent cache entry for 'jit__power'
      because it took < 1.00 seconds to compile (0.04s)

40 ms of XLA compilation, declined by the persistent cache for being too cheap to cache, to differentiate a scalar power. On short fits this is a real fraction of the run: ppatf2_phospho completes a 10-start trf fit in 30 s total.

It makes an optional extra effectively mandatory. _require_jax documents "An all-linear fit never reaches here" — but loguniform_var is the common way to declare a rate constant, so in practice most gradient fits are log-scaled and cannot run without pybnf[jax]. Of six published fitting jobs I ran (Erickson 2019 igf1r, Jaruszewicz-Blonska 2023 reduced_onoff, Kirsch 2020 p38atf2_binding and ppatf2_phospho, Lin 2021 nyc_multiphase, Salazar-Cavazos 2019 egfr_simpull), all six declare their parameters loguniform_var.

It also imports a heavyweight dependency into a path that does not otherwise need it. The ODE forward sensitivities come from bngsim's CVODES in C++; JAX is pulled in solely for this scalar factor.

Suggested fix

Give each scale an analytic d_from_sampling_space(u) alongside from_sampling_space_jax, and have _d_theta_d_u call it:

• log10: 10**u * ln(10)
• natural log: exp(u)
• linear: 1

Keep the jax.grad route as the fallback for any scale that does not supply one, so a custom scale still works and the pybnf[jax] extra stays genuinely optional. That removes the XLA compile from the hot path and drops the hard jax requirement for the common log-scaled fit.

A unit test comparing the analytic derivative against jax.grad for each built-in scale would pin the two together.

Environment

PyBNF 1.6.0, bngsim 0.11.35 (v0.11.35-77-g32395e4), Python 3.12, macOS 15.