EpiAwareADTools

Automatic-differentiation safety machinery for the EpiAware modelling stack.

Why EpiAwareADTools?

EpiAwareADTools is the EpiAware org's shared home for AD-safety machinery and AD workarounds. It is deliberately framed as fixes we host while we try to fix things upstream: every entry is documented with the upstream package or issue where it ideally belongs, and each is deleted once that upstream fix lands.

Three 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 on every backend.

• The AD-safe evaluation hooks (cdf_ad_safe, logcdf_ad_safe, ccdf_ad_safe, logccdf_ad_safe, and pdf_ad_safe) let a wrapper package overload CDF evaluation for its own component types, backed by analytic Gamma and Beta CDF derivatives that SpecialFunctions leaves unimplemented.

• Correct upstream ChainRulesCore rules lifted into a backend that lacks them, currently LogExpFunctions.xlogy/xlog1py under Mooncake, which otherwise returns a wrong shape-gradient for any Gamma log-density differentiated at shape == 1.

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

Getting started

See the documentation for a full walkthrough.

using EpiAwareADTools, Distributions

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

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

3.0

Related packages

• ConvolvedDistributions.jl, ComposedDistributions.jl, ModifiedDistributions.jl, LoweredDistributions.jl and CensoredDistributions.jl import these AD-safe hooks in their own source and overload them for their component types, so their densities differentiate on every supported backend.

• DistributionsInference.jl is the emerging fit-protocol layer across those packages, where the AD-safety this package provides is what makes gradient-based fitting work.

Where to learn more

• Want to get started running code? Check out the Getting started documentation.

• Want to understand the API? Check out our API reference.

• Want to contribute to EpiAwareADTools? Check the open issues and the Contributing section below.

• Want to see our code? Check out our GitHub Repository.

Getting help

For usage questions, ask on the Julia Discourse (the SciML or usage categories) or the epinowcast community forum, our home for epidemiological modelling questions. Please use GitHub issues for bug reports and feature requests only.

Part of the EpiAware ecosystem

EpiAwareADTools is part of EpiAware, a set of composable tools for infectious disease modelling. See the other packages in the ecosystem.

Contributing

We welcome contributions and new contributors! Please open an issue or pull request on GitHub. This package follows ColPrac and is formatted with Runic.

How to cite

If you use EpiAwareADTools in your work, please cite it. Citation metadata lives in CITATION.cff, which GitHub renders as a "Cite this repository" button on the repository page.

Code of conduct

Please note that the EpiAwareADTools project is released with a Contributor Code of Conduct. By contributing, you agree to abide by its terms.