perf(ode-grad): hoist state-Jacobian eval out of per-parameter loop (gh#400)#401
Merged
Conversation
…gh#400)
In `sensitivity_derivs`, the state-Jacobian term ∂rate_r/∂x_j
(`eval_deriv_entry`) was re-evaluated inside the `for p in 0..d` loop
over estimated parameters, but its value is independent of the
parameter column `p` — only the `S[j,p]` multiplier varies. Evaluate
each transition's J_x entries once during the first column and reuse
them for the remaining columns, cutting the J_x expression-eval work
from P·nnz to nnz per transition per RK4 stage. The FOI-type rates
(dense J_x, ~8-9 nonzeros) benefit most.
Byte-neutral: column 0 keeps the exact original J_θ-then-J_x
evaluation order, so the shared binding cache is populated identically;
columns p>0 previously re-called `eval_deriv_entry` against the warm
cache and got exactly the values now stashed in `jx`. An earlier
fully-hoisted variant that evaluated all J_x before J_θ reordered
cache population and perturbed one gradient component by 1 ULP — the
column-0 interleave avoids that.
A/B gate (det_grad on the seir_observations FD-oracle fixture,
f64::to_bits):
baseline ll=c0598ed08da0c428
grad=c05a4bbe00ed4375 4070de40d6693ee5
bffbd780d1241668 405e2ba1ec971515
after identical — all four grad words + loglik bit-for-bit.
FD oracles green (3/3); sim --lib green (323/0).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Cuts the per-step cost of the ODE forward-sensitivity gradient (the NUTS-on-ODE wall for long-burn-in fits) by removing a redundant expression evaluation. First of the two cost-per-step levers in gh#400; the second (kill per-step allocation) is deferred — see below.
What
In
sensitivity_derivs, the state-Jacobian term∂rate_r/∂x_j(eval_deriv_entry) was re-evaluated inside thefor p in 0..dloop over estimated parameters, even though its value is independent of the parameter columnp— only theS[j,p]multiplier varies. This evaluates each transition'sJ_xentries once (during the first column) and reuses them across the remaining columns, cutting theJ_xexpression-eval work fromP·nnztonnzper transition per RK4 stage. The force-of-infection rates (denseJ_x, ~8–9 nonzeros) benefit most.Byte-neutral
The gradient is bit-identical before/after. The mechanism matters: column 0 keeps the exact original
J_θ-then-J_xevaluation order, so the shared binding cache is populated identically; columnsp>0previously re-calledeval_deriv_entryagainst the warm cache and got exactly the values now stashed. A fully-hoisted variant that evaluated allJ_xbeforeJ_θreorders cache population and perturbs one gradient component by 1 ULP — the column-0 interleave avoids that.A/B gate (
det_gradon theseir_observationsFD-oracle fixture,f64::to_bits):Verification
make testgreen (all phases: Rust 907, inference, OCaml, integration)Deferred: lever 2 (per-step allocation)
The companion optimization — threading a reusable
Rk4Workspacethroughrk4_stepto eliminate per-step heap allocation — was implemented and reverted. It is not byte-identical: pure storage relocation drifts one gradient component by 1–2 ULP (value trajectory provably unchanged; only the dense-binding sensitivity chain moves), revealing that the forward-sensitivity accumulation is ULP-fragile to RK4 buffer structure via FP-order sensitivity. Held pending a decision on whether to accept a numerically-equivalent (FD-oracle-validated) change or root-cause the fragility first. Tracked in gh#400.