ChainRulesCore extension

EpiAwareADToolsChainRulesCoreExt loads when ChainRulesCore is loaded alongside EpiAwareADTools. It is the extension the others build on. ReverseDiff and Mooncake both depend on ChainRulesCore, so loading either of those backends loads this extension as well.

What it registers

Two non-differentiability marks:

• @non_differentiable primal(::Any), so a reverse-mode backend that reads ChainRules stops at the tape-strip instead of tracing through it.

• @non_differentiable (nd::NonDifferentiable)(args...), the same mark for every instance of the nondifferentiable wrapper. It is registered once on the wrapper type rather than once per wrapped function. The wrapper's own body already strips its arguments and its result through primal, so this mark is a second line of defence for a ChainRules-consuming context beyond the four backends this package tests directly.

An rrule and an frule for each of the three internal CDF primitives:

• _gamma_cdf(k, θ, x) = P(k, x/θ), the gamma-CDF derivative.

• _gamma_logccdf(k, θ, x) = log(Q(k, x/θ)), its log-space survival companion.

• _beta_cdf(α, β, x) = I_x(α, β), the beta-CDF derivative.

Each rule reads the primal value and the three partials from the matching _value_and_partials helper in src/gamma_ad.jl or src/beta_ad.jl. The helpers, not the rules, hold the analytic formulas. The ForwardDiff Dual methods and the Enzyme rules call the same helpers, so no backend carries its own copy of the maths.

The frule matters as much as the rrule. Mooncake's forward mode is generated by lifting ChainRules, and a lift with no frule behind it calls ChainRulesCore.frule, gets nothing back, and fails with iterate(::Nothing).

What fails without it

Only these rules and their per-backend counterparts register a derivative for the incomplete gamma and beta functions. With the extension unloaded, a ChainRules-based backend differentiating a Gamma or Beta CDF reaches SpecialFunctions' own rules, whose shape partials are @not_implemented, and the gradient call errors. The ReverseDiff and Mooncake extensions lift the rules defined here, so without this extension they have nothing to lift.

Upstream target

SpecialFunctions.jl carrying complete ChainRules for gamma_inc (issue #531) and for beta_inc. The three pairs of rules are deleted when that lands, along with the internal primitives they cover. The two @non_differentiable marks outlive them, and go once a shared cross-backend stop-gradient primitive exists.