Internal Documentation

Documentation for EpiAwareADTools's internal interface.

Contents

• Internal Documentation
  • Contents
  • Index
  • Internal API

Index

• EpiAwareADTools.NonDifferentiable
• EpiAwareADTools._beta_cdf
• EpiAwareADTools._beta_cdf_value_and_partials
• EpiAwareADTools._dlogQ_da_tail_series
• EpiAwareADTools._gamma_cdf
• EpiAwareADTools._gamma_cdf_value_and_partials
• EpiAwareADTools._gamma_logQ
• EpiAwareADTools._gamma_logccdf
• EpiAwareADTools._gamma_logccdf_value_and_partials
• EpiAwareADTools._grad_p_a_maxiter
• EpiAwareADTools._grad_p_a_series
• EpiAwareADTools._rib_f
• EpiAwareADTools._t_ccdf
• EpiAwareADTools._t_cdf
• EpiAwareADTools._t_logccdf
• EpiAwareADTools._t_logcdf
• EpiAwareADTools._t_pdf_at_zero
• EpiAwareADTools._t_tail

Internal API

EpiAwareADTools.NonDifferentiable Type

struct NonDifferentiable{F}

A callable wrapper holding a function out of differentiation.

NonDifferentiable(f) (built via nondifferentiable; use that, not this type, as the public entry point) strips every argument to its primal via primal before calling f, and also strips f's RESULT the same way, so the call contributes no derivative on any supported backend regardless of what f itself does internally — see nondifferentiable for the full contract.

---

Fields

• f::Any

source

EpiAwareADTools._beta_cdf Function

_beta_cdf(α::Real, β::Real, x::Real) -> Any

AD-safe regularised incomplete beta I_x(α, β) — the Beta(α, β) CDF at x.

Primal goes through SpecialFunctions.beta_inc (same path the non-AD hot path uses). AD coverage is supplied by per-backend extensions:

• EpiAwareADToolsChainRulesCoreExt defines the reverse-mode rrule and forward-mode frule (analytical partials, primal via beta_inc).

• EpiAwareADToolsMooncakeExt lifts both the rrule and frule into Mooncake (reverse and forward mode).

• EpiAwareADToolsReverseDiffExt lifts the rrule into ReverseDiff.

• EpiAwareADToolsForwardDiffExt defines Dual methods on _beta_cdf directly (forward-mode dispatches on argument types, not via ChainRules).

• EpiAwareADToolsEnzymeExt supplies the direct Enzyme rule.

The α/β-partials that beta_inc's ChainRule leaves unimplemented are supplied by _rib_value_and_partials, following the continued fraction Boik & Robinson-Cox (1998) derive for the regularised incomplete beta's shape-parameter derivatives. This mirrors _gamma_cdf's role for gamma_inc (SpecialFunctions.jl issue #531) — there is no equivalent tracking issue open against beta_inc at the time of writing; this machinery is deleted if/when SpecialFunctions.jl grows one.

source

EpiAwareADTools._beta_cdf_value_and_partials Function

_beta_cdf_value_and_partials(
    α::Real,
    β::Real,
    x::Real
) -> NTuple{4, Any}

Primal value and analytical partials (Ω, dα, dβ, dx) for _beta_cdf. Shared by every per-backend AD extension so the formulas live in one place:

• dx = pdf(Beta(α, β), x)

• (dα, dβ) = _rib_value_and_partials(x, α, β)[2:3] directly when x <= α/(α+β); otherwise obtained via the reflection symmetry I_x(α,β) = 1 - I_{1-x}(β,α), so dα = -(∂I_{1-x}(β,α)/∂(3rd arg)), dβ = -(∂I_{1-x}(β,α)/∂(2nd arg)) (the derivation Boik & Robinson-Cox's own reflected-argument identity gives; see the module docstring above _rib_f).

The x <= 0/x >= 1 branches return the same constant (Ω, 0, 0, 0) / (Ω, 0, 0, 0) pair _beta_cdf uses for its primal-only path (Ω = 0 or 1), matching its early-return behaviour.

source

EpiAwareADTools._dlogQ_da_tail_series Function

_dlogQ_da_tail_series(a::Real, z::Real) -> Any

∂ log Q(a, z) / ∂a deep in the right tail, from the large-z asymptotic expansion of the unnormalised upper incomplete gamma, Γ(a, z) = z^{a-1} e^{-z} S with S = Σₙ ∏ⱼ₌₁ⁿ (a - j) / zⁿ, giving

S and its a-derivative S' are accumulated by the joint recurrence tₙ = tₙ₋₁ (a - n) / z, sₙ = (sₙ₋₁ (a - n) + tₙ₋₁) / z, which never divides by (a - j) and so is exact when a is an integer and the product terminates. The expansion is asymptotic, not convergent: its terms shrink until n ≈ z - a and then grow without bound, so the loop stops at that optimal-truncation point, leaving an error of roughly the last included term, ~e^{-(z-a)}. Both exits watch the combined magnitude |tₙ| + |sₙ|, not tₙ alone: at integer a the value series terminates (tₙ hits an exact 0) while the derivative series sₙ keeps contributing, and stopping on tₙ there silently truncates S'. Used by _gamma_logccdf_value_and_partials once Q < √eps(T) (~1.5e-8 in Float64): there z - a ≳ 16 for any shape, so the truncation error is at worst ~1e-7 and falls exponentially with depth, while the exact ∂P/∂a / Q quotient it replaces is losing a digit for every decade Q drops (measured against finite differences of the stock logccdf, both paths hold ~1e-6 relative error or better at the crossover, the series reaching ~1e-12 once Q < 1e-10). The iteration count to the eps exit grows like z - a ~ 5.6√a in the slow geometric regime (a large, z/a near 1), so the cap binds only above shape ~2e5, where accuracy degrades gradually rather than failing.

source

EpiAwareADTools._gamma_cdf Function

_gamma_cdf(k::Real, θ::Real, x::Real) -> Any

AD-safe Gamma CDF, P(k, x/θ).

Primal goes through SpecialFunctions.gamma_inc for every Real subtype it supports (Float64, Float32, BigFloat) — same path the non-AD hot path uses, full accuracy across all z/a regimes. AD coverage is supplied by per-backend extensions:

• EpiAwareADToolsChainRulesCoreExt defines the reverse-mode rrule and forward-mode frule (analytical partials, primal via gamma_inc).

• EpiAwareADToolsMooncakeExt lifts both the rrule and frule into Mooncake (reverse and forward mode).

• EpiAwareADToolsReverseDiffExt lifts the rrule into ReverseDiff.

• EpiAwareADToolsForwardDiffExt defines Dual methods on _gamma_cdf directly (forward-mode dispatches on argument types, not via ChainRules).

The α-partial that gamma_inc's ChainRule leaves as @not_implemented is supplied by _grad_p_a_series, following the series form Moore (1982) introduced as Algorithm AS 187 and that Stan (grad_reg_inc_gamma) and JAX (igamma_grad_a) both use. This whole machinery stands in for a differentiable gamma_inc upstream (SpecialFunctions.jl issue #531) and is deleted once that exists.

source

EpiAwareADTools._gamma_cdf_value_and_partials Function

_gamma_cdf_value_and_partials(
    k::Real,
    θ::Real,
    x::Real
) -> NTuple{4, Any}

Primal value and analytical partials (Ω, dk, dθ, dx) for _gamma_cdf. Shared by every per-backend AD extension so the formulas live in one place:

• dx = pdf(Gamma(k, θ), x)

• dθ = -(x/θ) · dx

• dk = _grad_p_a_series(k, x/θ)

The non-positive-x branch returns zeros for the primal and all three partials, matching _gamma_cdf's early-return behaviour.

source

EpiAwareADTools._gamma_logQ Function

_gamma_logQ(a::Real, z::Real) -> Tuple{Any, Any}

Accurate log(Q(a, z)), the log of the regularised UPPER incomplete gamma Q(a, z) = 1 - P(a, z), for _gamma_logccdf and its value-and-partials companion.

SpecialFunctions.gamma_inc returns (P, Q) from independent series/continued-fraction evaluations rather than by subtracting P from 1, so its own Q output stays a representable, accurate value far beyond where log1p(-P) underflows — P rounds to exactly 1 in the working float type well before Q itself underflows to 0. log(Q) is read directly from that output while it stays representable; once Q itself underflows, the log-space form loggamma(a, z) - loggamma(a) (the log of the unnormalised upper incomplete gamma over log Γ(a), itself finite far beyond that point) takes over. This mirrors the branch structure StatsFuns._gammalogccdf uses for the stock (non-differentiable) evaluator, so the two agree at implementation tolerance across the whole domain rather than only where log1p(-P) happens to hold up.

Returns (log(Q), Q), with Q taken straight from gamma_inc rather than recovered as exp(log(Q)) — exponentiating costs |log Q| * eps relative error and lands in the subnormals well before gamma_inc's own Q does — so _gamma_logccdf_value_and_partials can divide by the accurate survival.

source

EpiAwareADTools._gamma_logccdf Function

_gamma_logccdf(k::Real, θ::Real, x::Real) -> Any

AD-safe Gamma log survival, log(Q(k, x/θ)) — the analogue of _gamma_cdf for the log-space survival rather than the CDF.

Unlike the naive log1p(-_gamma_cdf(k, θ, x)), this never forms the CDF as a literal float and subtracts it from 1: _gamma_logQ reads the survival from SpecialFunctions.gamma_inc's own second output, computed independently of the CDF, so precision survives far into the right tail where the CDF itself has already rounded to exactly 1 (EpiAwareADTools#47). AD coverage follows the same per-backend pattern as _gamma_cdf: _gamma_logccdf_value_and_partials supplies the shared primal and partials the ChainRules rule, the ForwardDiff Dual methods, and the Enzyme rule all consume.

source

EpiAwareADTools._gamma_logccdf_value_and_partials Function

_gamma_logccdf_value_and_partials(
    k::Real,
    θ::Real,
    x::Real
) -> NTuple{4, Any}

Primal value and analytical partials (Ω, dk, dθ, dx) for _gamma_logccdf.

The partials of a log survival are relative derivatives, and for a Gamma they converge to finite, non-zero limits in the deep tail — -dx is the hazard rate, which tends to 1/θ — so they must stay accurate past the point where Q itself underflows. dx and dθ both reduce to the ratio f/Q and are formed fully in log space as exp(logpdf - Ω), well conditioned to arbitrary tail depth. dk divides _grad_p_a_series's ∂P/∂a by gamma_inc's own accurately-computed Q (returned alongside Ω by _gamma_logQ, never recovered as exp(Ω)) while Q ≥ √eps(T) — the scaling at which the series' absolute rounding error stays negligible against ∂Q/∂a — and hands over to _dlogQ_da_tail_series below that: the naive quotient loses relative accuracy long before Q underflows, flipping sign entirely by Q ≈ 1e-16 in Float64. The type-relative threshold keeps a BigFloat caller on the exact quotient to the depth its precision genuinely supports. At reduced precision (Float32) the two paths cross at a higher error floor, so dk carries ~1e-5 relative error in the transition band — a documented limit rather than a tunable. Every backend receives finite, accurate gradients across the whole tail (EpiAwareADTools#47).

The x <= 0 branch returns the same constant (0, 0, 0, 0) _gamma_cdf_value_and_partials uses for its primal-only path, since log(Q(k, 0)) = log(1) = 0.

source

EpiAwareADTools._grad_p_a_maxiter Function

_grad_p_a_maxiter(a::Real, z::Real, rtol::Real) -> Any

Term budget for _grad_p_a_series, sized to the series rather than fixed.

The terms z^n / Γ(a + n + 1) grow while n < z - a and decay after, by a factor z / (a + n) ≈ 1 - m / z at m terms past that peak, so the log of the term falls by about m^2 / (2z). Converging to a relative rtol therefore takes about

terms, which grows like √a along z ≈ a. Half again the decay estimate is carried as margin. The fixed 10,000 this replaces was exhausted above shape ~4e5 and truncated the sum silently, costing 5% relative accuracy at shape 1e6 and all of it by 1e7 (EpiAwareADTools#67).

The floor keeps every shape below ~4e5 on exactly the iteration sequence it had before, so nothing in the reachable range moves. The ceiling bounds the worst case at a few milliseconds and binds only above shape ~4e8. Non-finite inputs fall back to the floor rather than throwing on the Int conversion; the series' own z <= 0 guard and the underflow of term handle the degenerate arguments themselves.

source

EpiAwareADTools._grad_p_a_series Function

_grad_p_a_series(a::Real, z::Real; rtol, maxiter) -> Any

Partial of the regularised lower incomplete gamma with respect to the shape parameter — the term SpecialFunctions.gamma_inc leaves as @not_implemented in its ChainRule. Computed by term-by-term differentiation of the Tricomi absolutely-convergent series for P(a, z) = z^a e^{-z} Σ_{n ≥ 0} z^n / Γ(a + n + 1):

with ψ(a + n + 1) = ψ(a + n) + 1 / (a + n) propagated alongside the term recurrence term_{n+1} = term_n · z / (a + n + 1). Used by the reverse-mode rule in EpiAwareADToolsChainRulesCoreExt and by the forward-mode Dual methods in EpiAwareADToolsForwardDiffExt.

This is the analytic α-partial the upstream gamma_inc ChainRule should eventually supply: SpecialFunctions.jl issue #531 (implement the shape-parameter partial for gamma_inc as a convergent series) tracks it. The helper is deleted once that rule lands.

Accuracy

The two sums differ by far less than either is worth: log(z) P stays O(log z) while ∂P/∂a falls away as P saturates, reaching 1e-8 by survival Q ≈ 3e-7 at shape 1e5. The closing subtraction therefore discards most of the working precision, and the rounding of the Θ(√a) terms accumulates on top of that. Measured against a 512-bit evaluation of the same series, the relative error is

shape a	Q ≈ 0.5	Q ≈ 1e-3	Q ≈ 3e-7
1e3	4e-13	5e-11	5e-8
1e5	8e-12	5e-10	2e-6
1e7	8e-10	1e-7	1e-4

Absolute accuracy holds throughout — ∂P/∂a itself is O(1e-8) in the worst column — so a _gamma_cdf gradient is unaffected. The relative error matters only to _gamma_logccdf_value_and_partials, which divides by Q, and only in the band above its √eps switchover — a shape whose coefficient of variation is under 3% (EpiAwareADTools#67). What is left is a property of this series rather than of the implementation, and a tighter loop cannot reach it. The cheap route to it is not a uniform (Temme) expansion but _dlogQ_da_tail_series, which is already several orders of magnitude better than the quotient wherever z - a is large, well before the √eps switchover hands over to it.

maxiter defaults to _grad_p_a_maxiter rather than a fixed count so the series always runs to its own convergence. A fixed 10,000 silently truncated above shape ~4e5 — 5% relative error at shape 1e6 and complete loss by 1e7, since the term count grows like √a.

References

The series + digamma-recurrence form is Moore (1982), "Algorithm AS 187: Derivatives of the Incomplete Gamma Integral", Applied Statistics 31:330-335. The same construction is used by Stan (stan/math/prim/fun/grad_reg_inc_gamma.hpp) and JAX (jax._src.scipy.special.random_gamma_grad / igamma_grad_a) for the shape derivative of the regularised lower incomplete gamma.

source

EpiAwareADTools._rib_f Function

_rib_f(x::Real, p::Real, q::Real) -> Any

Continued-fraction coefficients and their p/q-partials for the regularised incomplete beta function I_x(p, q) — the shape-parameter derivatives SpecialFunctions.beta_inc leaves unimplemented, the beta_inc analogue of _grad_p_a_series for gamma_inc.

Follows the continued-fraction form Boik & Robinson-Cox (1998) derive for I_x(p, q) and its p/q partials (their equations for a_n, b_n and the corresponding ∂a_n/∂p, ∂a_n/∂q, ∂b_n/∂p, ∂b_n/∂q), evaluated via the standard modified Lentz-style two-term recurrence

with K = x^p (1-x)^{q-1} / (p B(p,q)), differentiated term-by-term so dA_n, dB_n accumulate alongside A_n, B_n in the same pass. Valid for x ≤ p/(p+q); the complementary regime uses the symmetry I_x(p,q) = 1 - I_{1-x}(q,p) (and correspondingly for the partials), applied in _beta_cdf_value_and_partials.

References

Boik, R. J., & Robinson-Cox, J. F. (1998). "Derivatives of the Incomplete Beta Function." Journal of Statistical Software, 3(1), 1-20. The recurrence structure (coefficients a_n/b_n and their partials) follows the reference C implementation in Caner Türkmen's betaincder (MIT licensed), ported here to Julia with a 1000-term default ceiling this package's own test suite validates against the paper's published table and against SpecialFunctions.beta_inc for widely disparate p/q near the x boundary, where fewer terms converge too slowly for full precision (issue #42; the upstream 100-term default left both the primal and the partials under-converged there).

source

EpiAwareADTools._t_ccdf Function

_t_ccdf(ν::Real, x::Real) -> Any

AD-safe Student-t survival  , the reflection of _t_cdf: the Student-t is symmetric, so the survival at x is the CDF at -x and inherits its small-tail accuracy unchanged.

source

EpiAwareADTools._t_cdf Function

_t_cdf(ν::Real, x::Real) -> Any

AD-safe Student-t CDF , composed over _beta_cdf.

The smaller of the two tails is always the one evaluated: below zero and   above it, with I the regularised incomplete beta _t_tail returns. Both the branch test and the x == 0 guard read primal rather than the argument itself, because ForwardDiff's == and iszero compare a Dual's partials as well and a seeded zero therefore tests unequal to 0.

At x == 0 the beta argument   sits at 1, where _beta_cdf's x-partial diverges while the inner derivative   is exactly zero, so the chain rule asks for 0 * Inf. Since   for every ν, the guard returns that constant plus a term linear in x carrying the true x-partial : x's primal is zero, so the value is untouched and the ν-partial stays the correct zero.

Arguments

• ν: degrees of freedom.

• x: the evaluation point.

source

EpiAwareADTools._t_logccdf Function

_t_logccdf(ν::Real, x::Real) -> Any

AD-safe Student-t log survival  , the reflection of _t_logcdf.

source

EpiAwareADTools._t_logcdf Function

_t_logcdf(ν::Real, x::Real) -> Any

AD-safe Student-t log CDF .

Below zero this is the log of the directly computed small tail, so it stays finite and accurate arbitrarily deep into the left tail. Above zero the CDF is at least and log1p(-tail) reads it off the same small tail, so neither side ever reconstructs a small number as a difference of two large ones. The x == 0 guard defers to _t_cdf for the reason set out there.

Arguments

• ν: degrees of freedom.

• x: the evaluation point.

source

EpiAwareADTools._t_pdf_at_zero Function

_t_pdf_at_zero(ν::Real) -> Any

The Student-t density at the origin,  .

Kept in the same beta parameterisation as _t_tail, and used only by _t_cdf's x == 0 guard, which needs the finite x-partial the incomplete-beta composition cannot supply there.

source

EpiAwareADTools._t_tail Function

_t_tail(ν::Real, x::Real) -> Any

The smaller Student-t tail   , read straight off the regularised incomplete beta as

Composed over _beta_cdf rather than carrying rules of its own, so every backend that covers _beta_cdf — the ChainRulesCore rrule/frule, the ReverseDiff and Mooncake lifts of those, the ForwardDiff Dual methods, and the direct Enzyme rule — differentiates straight through this composition.

The identity is even in x, so the same expression serves the lower tail at a negative x and the survival at a positive one. Reading the small tail directly, rather than reconstructing it as  , is what keeps _t_logcdf accurate deep into the tail: at ν = 5 and x = -1e8 the direct value agrees with the incomplete beta's own small-argument asymptote to nine significant figures, where the stock evaluator has already lost all but the first.

source