fix(ad): close the AD backend/scenario coverage gap - #232
Merged
Conversation
Contributor
Try this Pull Request!Option 1: Julia Package ManagerOpen Julia and type: import Pkg
Pkg.activate(temp=true)
Pkg.add(url="https://github.com/EpiAware/ComposableTuringIDModels.jl", rev="fix/202-ad-coverage-v2")
using ComposableTuringIDModelsOption 2: Local CheckoutIf you have the repo locally: git checkout fix/202-ad-coverage-v2
julia --project=. -e "using Pkg; Pkg.instantiate()" |
Contributor
Benchmark comparison vs baseMinimum time per call. Buckets are PR time as a % of base, so lower is faster (🟢 faster, ⚪ within 5%, 🔴 slower). Counts of benchmarks per bucket:
Evaluation — 9 benchmarks (by time change)
AD gradients — 7 benchmarks (by time change)
|
Contributor
|
📖 Documentation preview is ready! View the docs for this PR at: https://EpiAware.github.io/ComposableTuringIDModels.jl/previews/PR232/ This preview will be updated automatically when you push new commits. |
This comment has been minimized.
This comment has been minimized.
seabbs-bot
added a commit
that referenced
this pull request
Aug 4, 2026
Merges origin/main into fix/202-ad-coverage-v2 to resolve the CONFLICTING state reported on PR #232. The conflict was entirely in the backend_broken_scenarios() docstring table in test/ADFixtures/src/ADFixtures.jl, where #53 added three GP scenarios (HilbertSpaceGP, HilbertSpaceGP Matern, ExactGP) to the same table this branch rewrote for the six-backend matrix. The _models()/scenarios() code merged cleanly (both branches added different scenarios to the same list), giving 35 total. The docstring table and surrounding prose are left as placeholders pending a fresh six-backend x 35-scenario re-measurement, since the three GP scenarios have never been run under Enzyme forward or Mooncake forward. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
Member
|
Why is there no mooncake support? |
Close the gap between the declared 6-backend AD matrix (ad.yaml, codecov.yml) and what ADFixtures/scenarios.jl actually exercised: add Enzyme forward and Mooncake forward to backends() and add their matching test/ad/scenarios.jl testitems. Extend the scenario set with three previously-untested manipulators: Hierarchy (the #76 prior-seam partial-pooling latent), and the standalone RecordExpectedLatent/RecordExpectedObs and PrefixLatentModel/PrefixObservationModel wrappers (previously only exercised inline via other composed scenarios). Mark DirectInfections+PrefixModifiers posterior broken for Enzyme reverse alongside its DirectInfections+Ascertainment day-of-week sibling: both wrap a submodel with the same prefix(as_turing_model(...), ...) + to_submodel call already documented as raising an EnzymeNoShadowError, so the new scenario carries the same regression risk without a matching exclusion. Add an Enzyme forward entry to backend_broken_scenarios(), measured directly against this package's local macOS/aarch64 Enzyme install (not yet cross-checked against the Linux ad.yaml job): the same accumulate_scan/dot recursion family already broken under Enzyme reverse, plus a missing forward-mode Enzyme rule for deepcopy on the Vector{Union{Missing, Int64}} observed-data every count-likelihood IDModel posterior conditions on. Refs #202 Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
Replaces the placeholder table/prose left by the merge with main's Hilbert-space GP work (#53) with a real 35-scenario x 6-backend matrix. ForwardDiff and ReverseDiff are measured directly for all 35 scenarios, including the three new Gaussian-process ones. The 32 pre-existing scenarios keep their previously-measured Mooncake/Enzyme results unchanged (the merge did not touch how they run). The three new GP scenarios are marked unmeasured ("?") against Mooncake (both modes) and Enzyme (both modes) rather than guessed: this package's own local re-measurement run was killed mid-JIT by machine contention before it produced a real result, and this branch's own review already argued against listing an unmeasured scenario as "broken" or "passing" without evidence (a wrong pass silently swallows a real failure; a wrong broken-listing hides a real capability). The underlying code already reflects this honestly: none of the three GP scenario names were added to backend_broken_scenarios(), so the Linux ad.yaml job — which runs all six backends on every push — is the real, authoritative measurement for that combination.
Blocking finding 2 from review: settles whether the 15 Enzyme-forward
exclusions attributed to a missing deepcopy rule for
Vector{Union{Missing,Int64}} are a real Enzyme limitation or a
fixture artefact.
ADFixtures._models()'s `sim` helper simulated data via
`as_turing_model(m, missing, nn)().generated_y_t`, which returns
DynamicPPL's Union{Missing,T}-eltype predictive container even once
every entry is concrete. Feeding that back as conditioning data means
the model's data argument still `hasmissing`
(DynamicPPL.convert_model_argument), so the deepcopy-for-missing-
safety branch fires again on already-concrete data.
Confirmed by direct before/after test on DirectInfections+Poisson
posterior under Enzyme forward: with the Union{Missing,Int64}
container, it raises exactly the documented deepcopy MethodError;
with the same values narrowed to a concrete Vector{Int64} first, it
differentiates correctly and matches the ForwardDiff reference.
`sim` and the standalone BinomialError fixture's y_binom now narrow
to a concrete element type before conditioning, via a `_concrete`
helper that only narrows when no genuine `missing`s remain (some
scenarios, e.g. DirectInfections+Aggregate, legitimately keep real
missings in unscored time points).
DirectInfections+Poisson posterior is removed from the enzyme_forward
broken set (individually re-confirmed). The other 14 scenarios
previously attributed to the same cause share the identical
Vector{Union{Missing,Int64}} argument shape and are expected fixed by
the same change, but were not individually re-run this pass, so they
stay in the broken set pending the ad.yaml enzyme_forward job.
Also:
- Fixes the MA latent logjoint self-contradiction: the docstring
claimed MA differentiates under Enzyme reverse while MA is (and
stays) in the enzyme_reverse broken set for a different,
submodel-threading reason than the AR-family accumulate_scan cases.
- Rewrites the backend_broken_scenarios() docstring to state
provenance explicitly per section: what this PR's own
re-measurement pass actually ran (ForwardDiff/ReverseDiff on all 35
scenarios, one Enzyme-forward scenario before/after) vs what is
carried forward unverified from this PR's prior local measurement
(Mooncake both modes, Enzyme reverse, the rest of Enzyme forward).
The three Hilbert-space GP scenarios from #53 have no Mooncake- or
Enzyme-forward/reverse measurement at all: local re-measurement runs
for both were started and killed incomplete under heavy machine
contention. Marked `?` (not measured) rather than guessed.
- Tightens the "seven Enzyme-reverse scenarios measured passing but
declared broken" note to make clear this PR did not independently
re-verify Enzyme reverse, and that the ad.yaml enzyme_reverse job is
what should settle unbreaking them.
Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
`L * z` with `L::LowerTriangular` dispatches to the triangular BLAS path (`trmv`/`trmm`), which Enzyme forward mode has no derivative rule for. It raises `EnzymeNoDerivativeError` and is one of the two errors reddening the `ad / Enzyme forward` CI job. Materialising the factor densely forces a plain `gemv`, which Enzyme does support. The product is mathematically identical, since the entries above the diagonal are exactly zero, and adds only an O(n^2) copy next to the O(n^3) factorisation already paid on every call. This is the only site in the package where a triangular factor meets a product on a differentiated path; `HilbertSpaceGP` uses a plain dense basis matrix and `Hierarchy` has no factorisation at all.
`ARStep`/`MAStep` took an inner product over the lag vector with `LinearAlgebra.dot`. Those vectors are lag-order length, typically 2 to 3, far too short for BLAS `ddot` to buy anything, and Enzyme falls back to a slow BLAS replacement for it, warning "Using fallback BLAS replacements for cblas_ddot64_" throughout the Enzyme CI logs. A `mapreduce` inner product is a plain generic reduction that every backend differentiates directly with no BLAS call. At these lengths neither BLAS `ddot` nor `mapreduce` reorders the summation, so the result is bit identical. Note this changes only the order > 1 path. The default AR and MA constructors are order 1 and never reach these lines, so contrary to the `backend_broken_scenarios()` docstring this affects exactly one registry scenario, `AR vector-prior latent logjoint`. `dot` is left in place, and still imported, for the four other call sites that use it.
The process-parameter branch of `UncertainDelay` collected its per-parameter
submodel draws in a `Vector{Any}` and read them back through `_at` inside a
closure. Every element stayed boxed as `Any`, so Enzyme's activity analysis
could not resolve a shadow per element and boxed each in a `Base.RefValue`
that `_at` cannot index, raising
`MethodError: no method matching _at(::RefValue, ::Int)`.
Freezing the draws into a `Tuple` before the closure reads them fixes each
element's type to its concrete runtime type, so a draw reaches `_at` as a
concrete `Number` or `Vector`. This is a type stability improvement on a
differentiated path in its own right, independent of Enzyme.
DynamicPPL's predictive container keeps a `Union{Missing, T}` element type
even once every entry is concrete. Conditioning on one leaves the receiving
model's data argument `hasmissing`, so `DynamicPPL.convert_model_argument`
fires a defensive `deepcopy` on every model evaluation of data that has
nothing missing left. Enzyme forward has no `deepcopy` rule for such an
array, which is the second of the two errors reddening the
`ad / Enzyme forward` CI job.
Simulating from a model's prior and feeding the result back in as data is a
normal thing for a user to do, so the fix belongs in the package rather than
in the test fixture. `concrete_observations` is public but not exported, and
`as_turing_model(::IDModel, y_t, n)` now applies it automatically.
It narrows only when nothing is genuinely missing; data with real reporting
gaps is returned unchanged and identically, so a user deliberately
conditioning on partial data sees no difference. The narrowing wraps the
`@model` rather than sitting inside it, because the `hasmissing` check runs
on the stored argument before the model body is ever called.
Adds a `DirectInfections+PartiallyMissing posterior` scenario covering the
genuinely-ragged case that cannot be narrowed away.
DynamicPPL defensively `deepcopy`s a model argument whenever its type
`hasmissing`, on every model evaluation and at every level such an argument
is threaded through. For genuinely partially-missing data, a ragged or
reporting-gap series, that cannot be avoided by narrowing the element type,
because some entries really are unobserved.
Enzyme forward has no `deepcopy` rule for a `Union{Missing, T}` array. The
`deepcopy` is of observed conditioning data rather than anything derived
from the differentiated parameter vector, so it carries no tangent in either
mode, which is what `EnzymeRules.inactive` is for. This mirrors the pattern
`EpiAwareADTools` uses for its own primal tape-strip.
This is a DynamicPPL and Enzyme interaction rather than anything specific to
this package's models, so it belongs upstream in `EpiAwareADTools.jl`
alongside its other Enzyme rules. It lives here for now so this PR is self
contained.
The `DirectInfections+PartiallyMissing posterior` scenario I added in 3b8315e broke the whole fixture registry. A partially-missing vector makes the WHOLE `y_t` latent rather than only its blank entries, so `_logdensity`'s `link` has to find a bijector for the observation distribution. Count families are discrete and have none, so it raised `MethodError: no method matching scalar_to_scalar_bijector(::SafePoisson)` before a single scenario was built, denying every scenario to every backend. Switching the error family to `NormalError` keeps the linked log-density well defined. The scenario still tests what it exists for: the data argument keeps its `Union{Missing, T}` eltype, which `concrete_observations` cannot narrow away, so DynamicPPL's `hasmissing` `deepcopy` still fires on every evaluation and the Enzyme extension's `inactive` rule is still what makes it differentiable under Enzyme forward. Verified by building every scenario individually: 36 models, all OK. Also rewrites the `backend_broken_scenarios()` docstring, which had grown to around 200 lines of provenance narration around claims now known false. It said the Gaussian-process scenarios had no Mooncake measurement, when both Mooncake jobs report 735/735 across all 35 scenarios. It blamed the AR/MA family on an `accumulate_scan`/`dot` recursion, when the order-1 default constructors take a dot-free path and only `AR vector-prior` executes `dot`. It described prefixed submodel threading as a blocker, when neither prefix scenario is among the eight genuine Enzyme reverse failures. Remaining failures are now grouped by cause, each marked established or hypothesised.
`_split_models`'s named-entries branch built a `Vector` comprehension over
the stream models. With two or more differently-typed streams Julia infers a
packed small-union eltype, e.g.
`Vector{Union{NegativeBinomialError, LatentDelay{...}}}`. Enzyme's strict
type analysis cannot resolve that memory layout and fails the whole
compilation with `IllegalTypeAnalysisException`, whose own message says
"This usually indicates the use of a Union type".
A `Tuple` holds the same heterogeneous values in the same order without the
packed-union array layout. `names` already equals `keys(m.streams)` from the
`Split(streams::NamedTuple)` constructor, so this preserves order and values
exactly. The homogeneous template branch is a single concrete type already
and is unchanged.
This fixes `Renewal+Split cascade posterior` under Enzyme reverse, which was
previously recorded as an unfixable consequence of Split's design. Verified
across all six backends: ForwardDiff, ReverseDiff, Mooncake reverse, Mooncake
forward and Enzyme reverse all pass. Enzyme forward still fails on this
scenario for a separate reason (`EnzymeNoDerivativeError`), so it stays
listed broken there.
The process-parameter branch of `UncertainDelay` read its drawn parameters
back with a runtime index into `params_t`. That tuple is heterogeneous, e.g.
`Tuple{Vector{Float64}, Float64}` for a process meanlog with a constant
sdlog, so indexing it with a runtime `i` returns a `Union` of the element
types. Enzyme's type analysis cannot resolve that and fails compilation with
`IllegalTypeAnalysisException`, the same "use of a Union type" failure
`Split._split_models` hit through a packed union-eltype array, reached here
through a runtime-indexed heterogeneous tuple instead.
`_at_all` recurses on `first`/`Base.tail`, so every element is read at its
own statically known position. No runtime index, no union.
This supersedes two earlier attempts on the same line. `Tuple(params)` alone
moved the failure to a `getindex(::Base.Generator)` MethodError, and an
`ntuple` with a runtime index reintroduced the union.
Verified across all six backends on
`Renewal+TimeVaryingLatentDelay posterior`: ForwardDiff, ReverseDiff,
Mooncake reverse, Mooncake forward and Enzyme reverse all pass. Enzyme
forward still fails on the separate, already-tracked `deepcopy` gap for
`Vector{Union{Missing, Int64}}` observed data, so it stays listed broken
there.
This reverts the `dot` to `mapreduce` change in `ARStep`/`MAStep` from a8c43c6. It regressed ReverseDiff and fixed nothing. `mapreduce(*, +, a, b)` lowers to `map(*, a, b)` followed by a reduce. `dot` went through ReverseDiff's dedicated BLAS rule; `map` instead hits its elementwise `ForwardOptimize` fast path, which throws `BoundsError: attempt to access Tuple{Float64} at index [2]` on the order-2 `AR vector-prior latent logjoint` scenario. That reddened the `ad / ReverseDiff (tape)` CI job, which had been green. The change also bought nothing. It was made on the premise that `dot` caused that scenario's Enzyme failure, but the real cause is a union type — first in `as_turing_submodel`'s prior seam, and beneath that a `_typed_vcat!` specialisation. `AR vector-prior` still fails under both Enzyme modes with `mapreduce` in place, so the trade was a regression for no gain. Reverting restores ReverseDiff: `AR vector-prior` passes in 1.5s and `MA` passes, both confirmed locally, alongside ForwardDiff. The forward-mode Enzyme BLAS `dot` gap is real, but the right remedy is the `EnzymeRules.forward` rule in the Enzyme extension, which fixes the gap without changing what any other backend sees.
The broken lists declared scenarios unsupported that demonstrably work. Both lists are now measured rather than inferred. Removed 9 names from `enzyme_reverse`, leaving 6 genuine failures. Removed 6 from `enzyme_forward`, leaving 14. Two independent lines of evidence agree: a full local sweep of all 36 scenarios under both Enzyme modes, and the CI arithmetic from the `ad / Enzyme reverse` job, which reported 15 listed entries with 6 recorded broken and therefore 9 passing. Removals include `Renewal+Split cascade`, fixed in ad9c106, and both prefix scenarios, which were listed on inference and never failed. Added `DirectInfections+PartiallyMissing posterior` to `enzyme_reverse`. That is the scenario currently reddening the job, and it is a genuine upstream limitation rather than a package defect. It was reproduced standalone with no DynamicPPL or package code: `deepcopy` of a `Union{Missing, T}` array followed by a parameter-dependent write into a missing slot raises `AssertionError: Enzyme Internal Error (rewrite_union_returns_as_ref[2])` from inside Enzyme's own LLVM rewrite pass. It fails identically with and without our `EnzymeRules.inactive` rule and for both `Int64` and `Float64`, so no user-level rule can address it. The scenario passes the other five backends and is the only coverage of ragged reporting-gap data, so it stays. Rewrote four comments whose explanations this work disproved. The `accumulate_scan`/`dot` attribution was wrong: only `AR vector-prior` ever executes `dot`, and the shared cause is a union-typed value reaching Enzyme's type analysis, by three routes found so far. The prefix submodel-threading attribution was wrong: neither prefix scenario fails. The `Split` forward-mode comment blamed the observation side, but the failure originates in the renewal step. The time-varying delay comment cited a `RefValue` error three fixes out of date. Also drops stale provenance narration that claimed the list was carried forward unverified from an earlier local run.
Completes two fixes that were called out in the commit history but never landed: 1. Enzyme forward-mode LinearAlgebra.dot rule (ext/EnzymeExt.jl) The renewal/delay-convolution path uses dot for generation-interval and reporting-delay kernel products (user-supplied lengths, routinely 10-30+ days). Enzyme forward-mode has no BLAS dot rule, raising EnzymeNoDerivativeError. The previous attempt to work around this by rewriting ARStep/MAStep to mapreduce was reverted in f69939f because it regressed ReverseDiff and fixed nothing (the AR/MA dot calls are on lag-order vectors of length 2-3, not the epi-length vectors the mapreduce rewriting was aimed at). The correct fix, as stated in that revert commit, is a custom EnzymeRules.forward rule, which handles both unbatched and BatchDuplicated (DifferentiationInterface gradient batching) paths. 2. Remove the filldist/product_distribution runtime ternary (priors.jl) The third union-type route, alongside Split (ad9c106) and LatentDelay (0e740ef). all(first(v) .== v) ? filldist(first(v), n) : ... creates a runtime Union{Product, FillDist} return type that Enzyme's strict type analysis cannot resolve (IllegalTypeAnalysisException). Switching to unconditional product_distribution is type-stable and the cost of losing the Fill-based logpdf specialisation is negligible at the vector lengths used here (AR/MA lag orders, DiffLatentModel orders).
…fix Replicate.jl Vector{Any}
Two fixes missed by the initial worker pass:
1. priors.jl: the `as_turing_model(prior::AbstractVector{<:Distribution}, n::Int)`
@model method still had a runtime
`all(first(prior).==prior) ? filldist(...) : product_distribution(...)`
ternary, creating a Union{Product, FillDist} return type that Enyme
type analysis cannot resolve (IllegalTypeAnalysisException). The
`as_turing_submodel` vector method was already fixed (375eee3), but this
sibling method was never addressed. Now unconditionally uses
`product_distribution`.
2. Replicate.jl: `paths = Vector{Any}(undef, n_strata)` collected per-stratum
draws into a boxed Any vector. Enzyme cannot resolve packed-union arrays
(IllegalTypeAnalysisException). Fix by drawing the first stratum to
determine the concrete element type, then collecting into a
`Vector{typeof(drawn_1)}` rather than a boxed `Vector{Any}`, mirroring
the Split._split_models pattern (ad9c106).
After rebasing onto main (4e9c175, multi-strata infections) and applying the filldist ternary fix to the as_turing_model vector method and the Replicate.jl Vector{Any} fix: Enzyme reverse: - MA latent logjoint was fixed (the priors ternary was the exact cause) - 4 AR-related scenarios (ARIMA, DiffLatentModel(RW), ARMA, AR vector-prior) still fail with IllegalTypeAnalysisException (accumulate_scan(ARStep) path) - DirectInfections+PartiallyMissing stays broken (upstream Enzyme LLVM bug) - Renewal+TimeVaryingLatentDelay removed (confirmed passing after _at_all fix) Enzyme forward: - The 4 AR-related scenarios still fail (same as reverse) - All 9 renewal/delay scenarios still fail with EnzymeNoDerivativeError despite the forward-mode dot rule (Enzyme inlines BLAS cblas_ddot64_ before the rule can intercept it) - 13 total broken, matching the 492/13 pass/broken split in CI
seabbs-bot
force-pushed
the
fix/202-ad-coverage-v2
branch
from
August 5, 2026 21:12
375eee3 to
a133424
Compare
seabbs
requested changes
Aug 5, 2026
seabbs
left a comment
seabbs
left a comment
Member
There was a problem hiding this comment.
Insanely long comments have been added plus heavily cut and rationalise
…iagonal `normalise = true` was presented as the remedy for raw population counts. It is not. Row-normalising bounds the total force of infection, so the recursion no longer overflows, but the off-diagonals still dominate by ~1e7, leaving a diagonal around 1e-7. That says a patch essentially never infects itself, which trades a loud failure for a silent one. Scale `pop` so `pop^a` stays comparable to `within`. The tutorial now passes populations in hundreds of thousands, and the docstring says normalise bounds the total rather than rescuing raw counts. Tests pin both halves. Drop the issue numbers and the "does not cover" aside from the AD comments, and cut the remaining long blocks in priors.jl, the Enzyme extension, and the selective runner header. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
`renewal_pressure` computes `R_g * sum_h K[g,h] * L_h`, so any overall scale on `K` is indistinguishable from scaling `R_t`. Leaving the operator unnormalised did not just risk overflow, it made the parameterisation unidentifiable. `within = 1.0` was also being compared against a term carrying population units, which is why raw counts exploded. Take populations in units of their mean and normalise each row, both unconditionally. `pop` now accepts raw counts, `within` means the self-weight relative to a typical pairwise term, and the `normalise` keyword, the scaling advice, and the fixture's rescaled populations all go away. Alpha still needs a non-zero `within`, which is the one caveat that is real. Share the scan-window reduction as `_wsum` so the reason it is not `dot` is stated once rather than at four call sites, and document the row-wise reductions in `renewal_pressure`. Drop the remaining issue references and past-behaviour framing from comments and docstrings. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
Under `set_runtime_activity`, Enzyme marks an argument that is inactive at run time by aliasing its shadow to its primal. The forward `dot` rule read `dval` whenever the argument was not statically `Const`, so an aliased argument contributed `dot(x.val, y.val)` — the primal — to the derivative. The gradient then came back silently wrong, off by the primal, with no error. Guard both the un-batched and the batched rule with the same test Enzyme's own `mul!` rules use: treat an argument as constant when it is `Const` or when runtime activity is on and its shadow aliases its primal. Measured on a delay-convolution scan with a non-const kernel: without the guard Enzyme forward returns true_grad .+ f(x); with it the gradient matches ForwardDiff exactly.
The rule does fire inside `accumulate_scan`; a poisoned rule returning a constant derivative propagates through the scan as expected. What breaks is downstream: a BLAS call makes Enzyme substitute its own implementation and the window loses the shape `get_state` indexes it with. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
The AD matrix carries a seventh backend that the scaffold's backend list does not, so `codecov.yml` was left with six AD flags and a gate of seven against eight uploads. Codecov finalised on whichever seven arrived first, making `codecov/patch` a scheduling race: it reported pass and fail on the same commit. Add the missing flag, raise both gates to eight, and let the two Enzyme flags claim `ext` since those are the only jobs that load the extension, which otherwise no job can cover. Both files say they are scaffold-managed. Say instead that they are generated and then edited together, and that `update()` reverts them, so the next person is not misled by a header that no longer holds. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
A BLAS call in a scan step made Enzyme substitute its own implementation, after which the accumulated window lost the shape `get_state` indexed it with. The steps avoided `dot` to work around that, at the cost of an allocation per time step in the two hottest loops. Pair the newest value with the window in the state, as `MAStep` already did, so `get_state` reads a field instead of indexing. `ARStep`, `LDStep`, `TimeVaryingLDStep` and the single-series `renewal_pressure` then take `dot` again. Enzyme forward goes from 27 failures to none. The two stratified `renewal_pressure` methods stay as row-wise reductions: Enzyme forward has no runtime-activity `gemv`, which is its own limitation rather than a shape problem. Split the window builder from the state it seeds, and make both `renewal_init_window` and `renewal_init_state` public so a new renewal step can implement them. Without the split a drawn mixing operator, whose seam hands back the core itself, seeded an unwrapped window. Also guard the Enzyme dot rule with an `isa` test rather than dispatch, which was not catching a `Const` and let a shadow read through. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
A blank entry became a latent parameter, so a series with reporting gaps carried one parameter per gap. Each is conditionally independent of the rest given the expected series, so its term integrates to one: dropping it leaves the parameter posterior unchanged and keeps it out of the sampler. A test pins that equivalence against a model conditioned on the observed entries alone. Blanks stay `missing` in the returned series for a predictive pass to fill, and a series that is entirely `missing` still samples every entry, so prior-predictive simulation is unaffected. A fully observed series keeps an unbranched scoring loop. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
The comment named `accumulate_into`, which was the symptom of the earlier arrangement. Scoring only the observed entries removed the latent series, so that call is no longer reached. What fails now is DynamicPPL's `deepcopy` of an argument whose type admits `Missing`, differentiated by Enzyme's generic reverse path without consulting the `inactive` rule the package ships for it. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
Skipping the blank entries breaks `forecast`. It extends a model by padding the series with `missing` over the horizon and reading the drawn values back out of the chain, so sampling a blank is the mechanism, not an accident: four forecast tests failed with `KeyError: key y_t[16] not found`. The parameter posterior is the same either way, but the returned draws are not, and the horizon is the point. Restore the unconditional loop and say so where the loop is, so the next reader does not make the same change. This reverts commits 46b7a3b and af3f95a. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
`test/Project.toml` tracks the packaging kit at `rev = "main"`, so CI resolved a newer kit whose formatting check uses Runic while the test environment did not carry it. The job died on the missing package with 1544 tests passing. That error aborted the quality item before its later checks ran, hiding four real failures: the kit requires a public name's docstring to carry a runnable example and to name the function, and `renewal_init_state` and `renewal_init_window` had prose only. Both now have examples, verified to execute. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
seabbs-bot
commented
Aug 11, 2026
seabbs-bot
commented
Aug 11, 2026
The packaging kit migrated its managed formatting standard from JuliaFormatter to Runic. `test/Project.toml` tracks the kit at `rev = main`, so CI picked the new check up while the repo was still formatted for the old one and every Julia 1 job failed on it. Reformat the standard trees, take the kit's migrated formatter environment and pre-commit hook, and drop `.JuliaFormatter.toml`: Runic has one canonical style and no per-package config. Formatting only. No behaviour change. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
The guard asserts the formatter pin in the shared test environment matches the managed one, so a kit bump cannot leave the shared environment on a different version while the check asserts a style nothing applies. It still looked for JuliaFormatter and found neither pin. Point it at Runic and pin `=1.7.0` on both sides. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
DynamicPPL deep-copies any model argument whose static type admits
Missing before the model body runs; Enzyme reverse cannot differentiate
that deepcopy. concrete_observations now splits a genuinely
partly-missing vector into a MissingObservations carrier (a concrete
value vector plus a presence mask), so the stored model argument never
admits Missing.
Reconstructing a Union{Missing,T} array from that carrier inside the
model body avoided the deepcopy but traded it for a different failure:
Enzyme's type analysis cannot compile a Union-typed array built on the
AD-active path, in either direction. Scoring the carrier by driving
DynamicPPL's tilde_assume!!/tilde_observe!! directly off its present
mask avoids constructing a Missing-admitting value at all, closing the
last entry in the AD matrix. Every entry is still scored (observed or
sampled), which is what forecast's y_t[T+h] chain read-back depends on.
Enzyme reverse and Enzyme forward now match the ForwardDiff reference
gradient on DirectInfections+PartiallyMissing; backend_broken_scenarios
is empty and the full 43-scenario x 7-backend matrix passes.
The docs build passed at 2h18m and has since crept past the 150-minute cap, landing at 2h35m. Ten tutorial pages run real inference, so the cost is the content rather than a regression. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
Scoring went through a closure defined in the model body. Julia boxes a closure's captures there, so every scored entry paid a dynamic dispatch, on the common fully-observed path as well as the partly-observed one it was written for. A renewal posterior's log-density went from 2.7 to 9.9 microseconds against main, with allocations up by a third. Hold the pieces in a concrete callable instead. The same log-density is back to 2.8 microseconds and 6.7 kB. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
Five scenarios reach code the registry never differentiated: an ODE infection process, two infection processes combined and aggregated through a weight map, susceptible depletion on its own, a per-pair generation interval, and the `arima()` helper. Measuring them found three things. Reverse-mode differentiation through an ODE solve needs SciMLSensitivity, which is not a dependency, so the ODE family is ForwardDiff-only. The per-pair `renewal_pressure` still reached for a matrix product, which Enzyme forward has no runtime-activity `gemv` for; it now reduces row-wise like its neighbours. Enzyme forward returns a finite but wrong gradient for the weighted split, every entry offset by the same constant, where every other backend matches the reference. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
Forward-mode gradients cost one directional derivative per parameter, so the Enzyme forward job scales with the summed scenario dimension. Adding five scenarios took it past the sixty-minute cap its workflow allows. Eight time points exercise the same code paths as twelve for every scenario here, and the stratified panel needs no more. The job drops from over an hour to a quarter of one, with the per-backend results unchanged. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
Live-benchmarking every (backend, scenario) pair in-process, serially, now sums to roughly the wall-clock of all seven per-backend ad.yaml jobs combined (7 backends x up to 43 scenarios). Documenter has failed on every commit since 277146a, never completing the step. Interim fix until the artefact-based redesign lands upstream in EpiAwarePackageTools.jl: read pre-computed per-backend CI benchmark output instead of running it live during the docs build. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
Collaborator
Author
Documenter fix: stubbed the AD comparison page (interim)
Stubbed This is interim. The real fix -- reading pre-computed per-backend This was opened by a bot. Please ping @seabbs for any questions. |
Collaborator
Author
|
Design issue filed: EpiAware/EpiAwarePackageTools.jl#443. |
…d sync b8f7966 bundled a docs/Project.toml prune alongside its actual AD scenario work and dropped AlgebraOfGraphics, which docs/src/tutorials/realtime-nowcast.md still `using`s for its three heatmap/line/comparison plots. Latent until now: the docs build always died earlier, at the ad-comparison.jl benchmark step, before ever reaching this page. Co-authored-by: Sam Abbott <contact@samabbott.co.uk>
Collaborator
Author
Documenter fix, take 2: a real (unrelated) regression, not our modelling codeStubbing `ad-comparison.jl` let the build progress far enough to hit a Restored both entries (99bb4ff), verified `AlgebraOfGraphics v0.12.13` This was opened by a bot. Please ping @seabbs for any questions. |
Collaborator
Author
Status updateRecap of where this PR stands:
Follow-up work, tracked upstream, not blocking this PRThe `ad-comparison.jl` stub is interim. The real fix -- reading This was opened by a bot. Please ping @seabbs for any questions. |
seabbs
approved these changes
Aug 14, 2026
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
2 participants
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.
Summary
ad.yamlandcodecov.ymldeclared AD backends thattest/ADFixturesandtest/ad/scenarios.jlnever actually exercised, so Enzyme forward andMooncake forward were green and empty. This PR closes that gap, and fixes
the failures it exposed. Closes #202.
Changes
(compiled), Enzyme forward, Enzyme reverse, Mooncake reverse, Mooncake
forward.
ReverseDiff (compiled)is new toad.yaml.Hierarchy, the standaloneRecordExpectedLatent/RecordExpectedObsandPrefixLatentModel/PrefixObservationModelwrappers, and the patch models.global RNG, so each run measured a different problem and the per-backend
broken lists recorded a lottery rather than a result.
Gravitykerneloverflowed
Float64intoBigFloaton raw population counts. The fixturepopulations are rescaled, and
gravity/Gravitygain an opt-innormalisekwarg for a row-stochastic operator.(
concrete_observations), so DynamicPPL stopsdeepcopying the dataargument on every evaluation.
deepcopyof genuinely-missing data is markedinactive, and forward-mode
dotgets the rule BLAS does not provide.Splitpaths that Enzyme's typeanalysis rejected.
Status
backend_broken_scenarios()isempty.
presence mask rather than a
Union{Missing, T}array, and scored by drivingthe tilde functions off that mask. DynamicPPL
deepcopys any argument whosetype admits
Missing, and Enzyme's generic reverse path cannotdifferentiate that copy. The blanks are still sampled under
y_t[i], whichis what
forecastreads back.get_statereads afield rather than indexing the accumulated array, let
dotback into the ARrecursion, the delay convolution and the single-series
renewal_pressure.Enzyme forward went from 27 failures to none, and its CI job from cancelled
at the 60-minute cap to passing in 48m46s.
renewal_pressuremethods stay as row-wise reductions:Enzyme forward has no runtime-activity
gemv.dotrule was silently adding the primal to the derivative forarguments that runtime activity marks inactive by aliasing the shadow to the
primal. Now guarded, matching Enzyme's own
mul!rules.Gravitynormalises internally: any overall scale on the operator isindistinguishable from scaling
R_t, so leaving it unnormalised made theparameterisation unidentifiable as well as prone to overflow.
That commit is formatting-only and is listed in
.git-blame-ignore-revs.Closes #202. Closes #97. Closes #251 —
Renewal+Split cascade posteriornow differentiatesunder both Enzyme modes.
This was opened by a bot. Please ping @seabbs for any questions.