Public Documentation

Documentation for EpiAwareADTools's public interface.

EpiAwareADTools.EpiAwareADTools Module

EpiAwareADTools

The EpiAware org's shared home for automatic-differentiation safety machinery and AD workarounds. Every entry is a fix hosted here while the real fix is pursued upstream, so each is documented with the upstream package or issue it stands in for and deleted once that lands.

Five families make up the current surface. The tape-strip pair primal and primal_distribution reduce an AD-wrapped scalar or distribution to its underlying primal, keeping a non-differentiable hyperparameter (an integration window, a clamp location) off the AD path. The AD-safe evaluation hooks cdf_ad_safe, logcdf_ad_safe, ccdf_ad_safe, logccdf_ad_safe, and pdf_ad_safe are extension points a wrapper package overloads for its own component types; their Gamma methods route through an analytic gamma-CDF derivative that stands in for the differentiability SpecialFunctions.gamma_inc leaves unimplemented, and their Beta methods do the same for SpecialFunctions.beta_inc's missing shape-parameter derivatives. Their TDist methods, and those for a t's location-scale wrapper, compose over that same beta machinery. nondifferentiable generalises primal's own discipline to an arbitrary user-supplied function: a deliberate, user-facing opt-out from differentiation, never a hidden default. logsumexp_stream is the shared streaming log-sum-exp accumulator for an infinite series over an unbounded discrete support, stopping only after a run of consecutive terms has left the total unchanged rather than at the first negligible term. The reparameterisation-trick pair fixed_draw and ad_eltype is the mirror image of the tape-strip pair: fixed_draw pins a random draw as a constant realisation the parameters vary against, and ad_eltype resolves the type a parameter-dependent accumulator combined with such a draw should be seeded at, so its derivative is not silently severed.

Per-backend behaviour (ForwardDiff, ReverseDiff, Enzyme, Mooncake, ChainRulesCore) is supplied by package extensions loaded when each backend is present.

Examples

using EpiAwareADTools, Distributions

# AD-safe Gamma CDF, differentiable in shape/scale on every backend.
cdf_ad_safe(Gamma(2.0, 1.0), 3.0)

# Strip an AD wrapper back to its primal value.
primal(3.0)

source

Contents

• public api
  • Contents
  • Index
  • Public API

Index

• EpiAwareADTools.EpiAwareADTools
• EpiAwareADTools.ad_eltype
• EpiAwareADTools.ccdf_ad_safe
• EpiAwareADTools.cdf_ad_safe
• EpiAwareADTools.fixed_draw
• EpiAwareADTools.logccdf_ad_safe
• EpiAwareADTools.logcdf_ad_safe
• EpiAwareADTools.logsumexp_stream
• EpiAwareADTools.nondifferentiable
• EpiAwareADTools.pdf_ad_safe
• EpiAwareADTools.primal
• EpiAwareADTools.primal_distribution

Public API

EpiAwareADTools.ad_eltype Function

ad_eltype(x::Real) -> Type{<:Real}

Resolve the element type a parameter-dependent accumulator should be seeded at, so a derivative threaded through it is not silently severed.

Accumulating a quantity that depends on differentiated parameters — for example summing per-draw contributions against a set of fixed_draw realisations — needs its running total, or any buffer it is written into, to carry the parameters' own AD-wrapped type from the FIRST term onward. Seeding it with a plain Float64 zero, or preallocating a buffer at Float64, works only until a differentiated term is combined with it: ForwardDiff and ReverseDiff (tape mode) each thread their own wrapper type (a Dual, a TrackedReal) through the live computation, and writing that wrapper into a concretely-Float64-typed slot raises a MethodError rather than quietly losing the derivative.

ad_eltype(x) returns the type to seed with: typeof(x) for a plain Real (so zero(ad_eltype(3.0)) === 0.0 when nothing is differentiated), and the promote_type across a Tuple's or AbstractArray's elements otherwise, so a caller mixing a differentiated parameter alongside a fixed_draw constant in the same container still resolves to the wider, differentiated type. Enzyme and Mooncake trace plain Float64 code by source transformation rather than threading a wrapper type through the primal computation, so ad_eltype is a harmless no-op — always Float64 — on those two backends; it does the real work on ForwardDiff and ReverseDiff, where the parameter genuinely carries a different runtime type.

Arguments

• x: the differentiated parameter(s): a Real, or a Tuple/ AbstractArray of such (mixed types promoted to their common type).

Examples

using EpiAwareADTools

T = ad_eltype(3.0)
total = zero(T)
for zi in fixed_draw.([0.1, 0.4, 0.9])
    total += zi
end
total

See also

• fixed_draw: the matching realisation-pinning half

• primal: the tape-strip discipline fixed_draw reuses

source

EpiAwareADTools.ccdf_ad_safe Function

ccdf_ad_safe(dist, u::Real) -> Any

AD-safe ccdf(dist, u): the survival  .

ccdf_ad_safe is the survival companion to cdf_ad_safe. Generic dispatch falls through to Distributions.ccdf; the Gamma method routes through the AD-safe so the survival differentiates w.r.t. the Gamma shape/scale.

An extension point: a downstream package adds methods for its own component types, the same pattern as pdf_ad_safe.

Arguments

• dist: the distribution whose survival is evaluated. The fallback carries no type bound, so a leaf implementing the Distributions generic without subtyping UnivariateDistribution is accepted.

• u: the evaluation point.

Examples

using EpiAwareADTools, Distributions

ccdf_ad_safe(Gamma(2.0, 1.0), 3.0)

source

EpiAwareADTools.cdf_ad_safe Function

cdf_ad_safe(dist, u::Real) -> Any

AD-safe cdf(dist, u) companion to logcdf_ad_safe.

Same dispatch idea: route Gamma through _gamma_cdf so a CDF evaluation remains differentiable under reverse-mode AD in its shape/scale. A downstream numeric kernel that evaluates components through this hook can add a method for a component type with a non-AD-safe cdf.

An extension point: a wrapper package adds methods the same way as pdf_ad_safe.

Arguments

• dist: the distribution whose CDF is evaluated. The fallback carries no type bound, so a leaf implementing the Distributions generic without subtyping UnivariateDistribution is accepted.

• u: the evaluation point.

Examples

using EpiAwareADTools, Distributions

cdf_ad_safe(Gamma(2.0, 1.0), 3.0)

source

EpiAwareADTools.fixed_draw Function

fixed_draw(x) -> Any

Pin a draw as a fixed, non-differentiated realisation for the reparameterisation trick.

Differentiating a quantity built from a random draw only makes sense when the draw is held fixed while the parameters vary against it — the common-random-numbers discipline. fixed_draw(x) strips any AD wrapper from x via primal: the identical mechanism primal already applies to a non-differentiable hyperparameter, given a distinct name here because the value being frozen is a REALISED DRAW, not a structural hyperparameter. It accepts exactly what primal does — a plain Real (returned unchanged, keeping its own float type), or a Tuple/ AbstractArray of such, stripped elementwise and recursively — so a single draw or a whole batch of them pins the same way.

CAUTION — fixed_draw only guards against an AD wrapper that survives the sampling step itself (for example a distribution whose sampler reads its Dual-valued bounds and returns a Dual). It does NOT by itself implement the common-random-numbers discipline: the draw must already be independent of the differentiated parameters — generated from their primal values, or before the parameters vary at all — or the resulting gradient is wrong rather than merely imprecise. See ad_eltype for the matching half: the type a parameter-dependent accumulator combined with a fixed draw should be seeded at.

Arguments

• x: the draw, or a container of draws, to pin as a constant.

Examples

using EpiAwareADTools

z = fixed_draw(rand(3))
length(z)

See also

• primal: the mirror-image tape-strip this delegates to

• ad_eltype: the matching accumulator element-type resolver

source

EpiAwareADTools.logccdf_ad_safe Function

logccdf_ad_safe(dist, u::Real) -> Any

AD-safe logccdf(dist, u): the log survival  .

logccdf_ad_safe is the log-survival companion to logcdf_ad_safe. Generic dispatch falls through to Distributions.logccdf; the Gamma method routes through _gamma_logccdf, which computes the survival directly rather than as  , so a survival term differentiates w.r.t. the Gamma shape/scale (the stock logccdf(::Gamma) calls _gammalogccdf, which has no ForwardDiff.Dual shape method and errors) and stays accurate far into the right tail, where itself has already rounded to 1 (EpiAwareADTools#47).

An extension point: a downstream package adds methods for its own component types, the same pattern as pdf_ad_safe.

Arguments

• dist: the distribution whose log survival is evaluated. The fallback carries no type bound, so a leaf implementing the Distributions generic without subtyping UnivariateDistribution is accepted.

• u: the evaluation point.

Examples

using EpiAwareADTools, Distributions

logccdf_ad_safe(Gamma(2.0, 1.0), 3.0)

source

EpiAwareADTools.logcdf_ad_safe Function

logcdf_ad_safe(dist, u::Real) -> Any

AD-safe logcdf(dist, u) for use inside differentiable integrands.

logcdf_ad_safe is the log-CDF member of the AD-safe hook family. Generic dispatch falls through to Distributions.logcdf. The Gamma method routes through _gamma_cdf so its ChainRulesCore.rrule is picked up by reverse-mode AD; without this, the integrand calls gamma_inc and breaks under every supported AD backend.

An extension point: a downstream package adds methods for component types whose stock logcdf is not AD-safe, the same pattern as pdf_ad_safe.

Arguments

• dist: the distribution whose log CDF is evaluated. The fallback carries no type bound, so a leaf implementing the Distributions generic without subtyping UnivariateDistribution is accepted.

• u: the evaluation point.

Examples

using EpiAwareADTools, Distributions

logcdf_ad_safe(Gamma(2.0, 1.0), 3.0)

source

EpiAwareADTools.logsumexp_stream Function

logsumexp_stream(
    log_term;
    atol,
    min_stable_terms,
    max_terms,
    strict
) -> NamedTuple{(:value, :terms_used, :converged), <:Tuple{Any, Int64, Bool}}

A differentiable streaming log-sum-exp accumulator over an unbounded discrete support.

logsumexp_stream(log_term) computes log(Σ_{k≥0} exp(log_term(k))) without materialising the full term sequence: it accumulates with the classic running-maximum/rescale log-sum-exp identity, updated incrementally as each log_term(k) arrives for k = 0, 1, 2, …. It stops only once the running total has been unchanged within atol for min_stable_terms CONSECUTIVE further terms — not at the first negligible term, which would silently truncate a heavy tail that dips low for a term or two and then recovers.

This is still a fixed-lookahead rule, not a proof of convergence: a sequence with a genuine plateau of min_stable_terms or more consecutive negligible terms, followed by a later resurgence, is indistinguishable from true convergence and will stop early regardless of what comes after. A caller expecting a distribution with a long flat middle should raise min_stable_terms accordingly.

Returns a NamedTuple (value, terms_used, converged): value is the accumulated log-sum (differentiable in whatever log_term closes over), terms_used is how many terms were consumed, and converged is true unless max_terms was reached first. When max_terms is reached without stabilising, the default (strict = true) raises a descriptive error rather than silently returning a partial sum; pass strict = false to receive the partial result with converged = false instead (the function itself never logs — a @warn/@info anywhere in this body would break Mooncake's whole-function transform even on an unreached branch, so a caller that wants a warning checks result.converged after the call returns and logs it there).

Every supported AD backend (ForwardDiff, ReverseDiff, Enzyme, Mooncake, ChainRulesCore) differentiates straight through value: the accumulator is plain generic Julia control flow over whatever real (or AD-wrapped) values log_term returns, calling no non-differentiable primitive, so it needs no bespoke per-backend rule — the same as a plain sum(logpdf.(...)) loop.

Arguments

• log_term: a function k::Int -> Real returning the log of the k-th term (k = 0, 1, 2, …), possibly carrying a live AD wrapper.

Keyword Arguments

• atol: the per-term change below which a term counts as "stable" (default 1e-12).

• min_stable_terms: how many CONSECUTIVE further terms must each leave the running total unchanged within atol before stopping (default 8).

• max_terms: the hard cap on terms consumed (default 10_000).

• strict: when true (default), raise an error if max_terms is reached without stabilising; when false, return the partial result with converged = false instead.

Examples

using EpiAwareADTools

# Σ_{k≥0} exp(-k), i.e. a geometric series with ratio exp(-1): the exact
# closed form is 1 / (1 - exp(-1)), so log of that is the reference value.
result = logsumexp_stream(k -> -Float64(k))
result.value ≈ log(1 / (1 - exp(-1)))
result.converged

source

EpiAwareADTools.nondifferentiable Function

nondifferentiable(f) -> EpiAwareADTools.NonDifferentiable

Hold a function out of differentiation: a deliberate, user-facing opt-out.

nondifferentiable(f) returns a callable that strips every argument to its primal via primal, calls f on the stripped arguments, and strips the RESULT the same way — so the call contributes exactly zero derivative on every supported backend (ForwardDiff, ReverseDiff, Enzyme, Mooncake, ChainRulesCore-consuming backends), regardless of what f computes internally. This generalises the discipline primal already applies to itself (the per-backend @non_differentiable/inactive/@zero_derivative rules) to an arbitrary function the caller names, written once here rather than threaded by hand through downstream code.

Because differentiation stays the norm and this is an explicit opt-out, every argument AND f's result must be something primal can strip — a Real, a Tuple or AbstractArray of such (nested to any depth), or a type a primal method already covers. Anything else raises a MethodError naming the missing primal method, a loud failure rather than a silently wrong derivative; add a primal method for that type (the same pattern primal_distribution follows for a distribution's parameters) to cover it. A STRUCT's constructor is itself callable, so wrapping it the same way — nondifferentiable(QuadratureGrid) — holds construction out of differentiation too, once the struct's own type has a primal method (there is deliberately no generic reflection-based primal fallback for an arbitrary struct: isstructtype is true for Dict, Module and every concrete function type — a closure, or typeof(sin) — as well as a user's own type, so a blanket fallback over every struct type would silently mishandle values this package was never asked to touch).

CAUTION — a captured value, not just an explicit argument, is also held constant: if f is a closure over a live differentiated value (rather than receiving it as an argument), that captured contribution is silently dropped too, consistently across every backend (confirmed directly: a closure capturing a component of the vector being differentiated reports the SAME reduced gradient under ForwardDiff, ReverseDiff, Enzyme, and Mooncake alike). This is the correct, deliberate consequence of "everything in this call is a constant" — never close over a value you still want differentiated.

Arguments

• f: the function (or callable, including a struct's own constructor) to hold out of differentiation.

Examples

using EpiAwareADTools

# A structural quantity that should stay fixed while parameters vary during
# optimisation or sampling: the midpoint of a quadrature window is *where*
# to evaluate, not something to estimate, even when `lo`/`hi` themselves
# carry a live AD wrapper elsewhere in the same computation.
window_midpoint(lo, hi) = (lo + hi) / 2
frozen_midpoint = nondifferentiable(window_midpoint)

frozen_midpoint(0.0, 1.0)

source

EpiAwareADTools.pdf_ad_safe Function

pdf_ad_safe(dist, t::Real) -> Any

AD-safe pdf(dist, t) for a component density inside a differentiable quadrature.

pdf_ad_safe is the density companion to cdf_ad_safe. Generic dispatch falls through to Distributions.pdf, and a downstream extension adds a method for a component whose stock pdf routes through functions that are not differentiable under the supported AD backends.

An extension point: a wrapper package hooks it so its modified components stay differentiable inside the quadrature, the same pattern as ccdf_ad_safe.

Arguments

• dist: the component distribution whose density is evaluated. The fallback carries no type bound, so a leaf implementing the Distributions generic without subtyping UnivariateDistribution is accepted.

• t: the evaluation point.

Examples

using EpiAwareADTools, Distributions

pdf_ad_safe(Gamma(2.0, 1.0), 3.0)

source

EpiAwareADTools.primal Function

primal(x::Real) -> Any

Strip any AD wrapper (ForwardDiff Dual, ReverseDiff TrackedReal, Enzyme/Mooncake duals) from a scalar, returning its underlying primal value.

The generic method is the identity on a plain real, so a non-AD call keeps the value's own float type (e.g. Float32). The per-backend extensions add the unwrapping methods: EpiAwareADToolsForwardDiffExt and EpiAwareADToolsReverseDiffExt supply value-reading methods, while EpiAwareADToolsChainRulesCoreExt marks primal @non_differentiable, EpiAwareADToolsEnzymeExt marks it EnzymeRules.inactive, and EpiAwareADToolsMooncakeExt gives it a @zero_derivative rule. Together these keep a non-differentiable hyperparameter — for example a quadrature window endpoint, which is just where to integrate — off the AD path on every backend.

Containers strip elementwise and recursively, so a nesting of them bottoms out at a scalar method. Tuple covers the nested per-component parameter tuples a composite distribution returns from params, which is what lets primal_distribution, and any caller mapping primal over params(d), handle a component whose parameter is itself a tuple. AbstractArray covers a vector-valued hyperparameter such as a grid of integration nodes, and returns a new container rather than stripping in place. Nothing passes through unchanged, so a wrapper distribution's absent bound — the nothing a one-sided truncated/censored stores — strips alongside its numeric siblings.

This is the sanctioned replacement for the underscore-prefixed _primal that ConvolvedDistributions.jl and CensoredDistributions.jl each carry internally; it stays hosted here until those packages depend on it directly.

Arguments

• x: the value to strip; a plain real is returned unchanged, a tuple or array is stripped elementwise into a new container of the same shape.

Examples

using EpiAwareADTools

primal(3.0), primal(((1.0, 2.0), 3.0)), primal([1.0, 2.0])

source

EpiAwareADTools.primal_distribution Function

primal_distribution(d) -> Any

Rebuild a distribution with its parameters stripped to their primal values via the type's positional constructor.

params(d) round-trips through the constructor for the Distributions.jl families used here, so mapping primal over the parameters and calling the wrapper constructor reconstructs the distribution with plain (AD-stripped) parameters. The check_args = false keyword is intentionally not passed: the original distribution already validated its parameters and the primal copy uses the identical values. Reconstructing from primal parameters is what lets a downstream package evaluate a non-differentiable quantity (such as a quadrature-window quantile) on an integration component without live Dual or TrackedReal parameters flowing into it.

A wrapper distribution whose params flattens its inner distribution's parameters and appends its own — Truncated and Censored — cannot round-trip through the positional constructor, so each has its own method rebuilding the inner distribution recursively and re-applying the primal bounds through the public truncated/censored. Any other distribution whose params does not match its constructor raises an ArgumentError naming the type rather than a MethodError naming a constructor the caller never wrote.

There is no type bound on d, so a leaf implementing the univariate interface without subtyping UnivariateDistribution rebuilds too, provided its params round-trips through its own constructor. One that defines no Distributions.params at all raises the same kind of ArgumentError.

Arguments

• d: the univariate distribution to rebuild from primal parameters.

Examples

using EpiAwareADTools, Distributions

primal_distribution(Gamma(2.0, 1.0))
primal_distribution(truncated(Gamma(2.0, 1.0), 0.0, 10.0))

source