diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
new file mode 100644
index 0000000..829f21c
--- /dev/null
+++ b/.github/workflows/docs.yml
@@ -0,0 +1,60 @@
+name: Docs
+
+on:
+ push:
+ branches: [main]
+ workflow_dispatch:
+
+permissions:
+ contents: read
+ pages: write
+ id-token: write
+
+concurrency:
+ group: pages
+ cancel-in-progress: false
+
+jobs:
+ build:
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v4
+
+ - uses: actions/setup-python@v5
+ with:
+ python-version: '3.11'
+ cache: pip
+ cache-dependency-path: docs/requirements.txt
+
+ - name: Install pandoc
+ run: sudo apt-get update && sudo apt-get install -y pandoc
+
+ - name: Install Python deps
+ run: pip install -r docs/requirements.txt
+
+ - name: Generate tutorial pages from .lhs
+ run: |
+ mkdir -p docs/mkdocs/docs/examples
+ for f in QueueModel OptionPricing GeneTranscription; do
+ pandoc --from=markdown+lhs --to=commonmark "tutorials/${f}.lhs" \
+ -o "docs/mkdocs/docs/examples/${f}.md"
+ done
+
+ - name: Build mkdocs site
+ run: cd docs/mkdocs && mkdocs build
+
+ - uses: actions/configure-pages@v5
+
+ - uses: actions/upload-pages-artifact@v3
+ with:
+ path: docs/mkdocs/site
+
+ deploy:
+ needs: build
+ runs-on: ubuntu-latest
+ environment:
+ name: github-pages
+ url: ${{ steps.deployment.outputs.page_url }}
+ steps:
+ - id: deployment
+ uses: actions/deploy-pages@v4
diff --git a/.gitignore b/.gitignore
index 656f417..7dea031 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,6 +11,10 @@ stack.yaml.lock
# Local dev config (kept locally, not shipped)
.hlint.yaml
+# mkdocs build outputs
+docs/mkdocs/site/
+docs/mkdocs/docs/examples/*.md
+
# Python
docs/.venv/
__pycache__/
diff --git a/README.md b/README.md
index 16e9fc0..34890ce 100644
--- a/README.md
+++ b/README.md
@@ -1,16 +1,19 @@
-
- 
- 
-
+
+
+
+
+ 
+
+
+# GradInf: Gradient Estimation as Probabilistic Inference
+
+GradInf is a research package that provides the Haskell implementation
+of *gradient inference*, a new approach to gradient
+estimation described in
+[this PLDI'26 paper](https://dl.acm.org/doi/abs/10.1145/3808321).
+
+## Overview
+
+
+ 
+
+
+GradInf automatically synthesizes gradient estimators
+(i.e., estimators of gradients of expected values [1]) for
+probabilistic programs. Users define a probabilistic
+program using a mix of ordinary Haskell and library-provided
+primitives. They then call `gradInfAD`, specifying their desired
+probabilistic inference strategy. The result is an unbiased estimate of
+the gradient of the expectation
+of the original program with respect to its parameter.
+
+### Starter Example
+
+We can write a simple queuing model:
+```haskell
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE RebindableSyntax #-}
+
+import Prelude hiding (flip)
+import Numeric.GradInf.Primitives.DeterministicPrimitives
+import Numeric.GradInf.Primitives.FlipCRN
+import Numeric.GradInf.Primitives.IterateP
+
+queueKernel :: forall m d i b mat.
+ (DeterministicPrimitives d i b mat, FlipCRN m d b)
+ => d -> i -> m i
+queueKernel theta x = do
+ let p = theta / (theta + if (isGreater x 25) :: b then 25 else toDouble x)
+ b :: b <- flipCRN p
+ let x' = if b then x + 1 else x - 1
+ return x'
+
+queueModel :: forall m d i b mat.
+ (DeterministicPrimitives d i b mat, FlipCRN m d b, IterateP m i)
+ => Int -> d -> m d
+queueModel n theta = do
+ x <- iterateP (queueKernel theta) 0 !! n
+ return (toDouble x)
+```
+The annotated probability distributions (e.g. `flipCRN`) specify
+*factorized coupling strategies*, which GradInf uses to form
+lower variance gradient estimators.
+We can now differentiate the program using the GradInf high-level API:
+```haskell
+import Data.Functor.Identity
+import Numeric.GradInf
+
+let thetaToDifferentiateAt = Identity 15.0
+let n = 50
+gradientEstimate <- sampler $ gradInfAD (queueModel n . runIdentity) (stratifiedImportanceResamplingInferenceAlg 1) thetaToDifferentiateAt
+print gradientEstimate
+```
+
+Here, `stratifiedImportanceResamplingInferenceAlg 1` is an *inference strategy*
+which enables sound and efficient gradient estimation.
+See the tutorials in this documentation site for more details on the API.
+
+ [1] *Monte Carlo Gradient Estimation in Machine Learning*, Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. *Journal of Machine Learning Research* 21, 132 (2020),
+
+## Index
+
+For more usage examples, check out the tutorial pages on applying GradInf to [queuing models](examples/QueueModel.md), [financial models](examples/OptionPricing.md), and [biochemical models](examples/GeneTranscription.md).
+*API documentation and additional tutorials will be added soon!*
+
+## Citation
+
+```bibtex
+@article{arya2026gradinf,
+title = {GradInf: Gradient Estimation as Probabilistic Inference},
+author = {Arya, Gaurav and Huot, Mathieu and Schauer, Moritz and Lew, Alexander K. and Saad, Feras A.},
+journal = {Proceedings of the ACM on Programming Languages},
+volume = {10},
+number = {PLDI},
+articleno = {243},
+month = jun,
+pages = {1864--1890},
+year = {2026},
+publisher = {Association for Computing Machinery},
+address = {New York, NY, USA},
+doi = {10.1145/3808321},
+}
+```
diff --git a/docs/mkdocs/docs/overrides.css b/docs/mkdocs/docs/overrides.css
new file mode 100644
index 0000000..29fa732
--- /dev/null
+++ b/docs/mkdocs/docs/overrides.css
@@ -0,0 +1,22 @@
+p {
+ margin: 0 0 18px;
+}
+
+.md-typeset h1, .md-typeset h2, .md-typeset h3, .md-typeset h4, .md-typeset h5, .md-typeset h6 {
+ margin-bottom: 18px;
+ color: black;
+ font-weight: 500;
+}
+
+.md-typeset .output {
+ color: black;
+ background: white;
+}
+
+/* Force the inline index-page logo to its `height` attribute. Material's
+ `.md-typeset img { height: auto }` otherwise overrides it. */
+.md-typeset img[alt="GradInf Logo"] {
+ height: 80px !important;
+ width: auto !important;
+}
+
diff --git a/docs/mkdocs/docs/overrides.js b/docs/mkdocs/docs/overrides.js
new file mode 100644
index 0000000..3c0cdf3
--- /dev/null
+++ b/docs/mkdocs/docs/overrides.js
@@ -0,0 +1,2 @@
+hljs.configure({languages:[]});
+hljs.initHighlightingOnLoad();
\ No newline at end of file
diff --git a/docs/mkdocs/mkdocs.yml b/docs/mkdocs/mkdocs.yml
new file mode 100644
index 0000000..c9ed84d
--- /dev/null
+++ b/docs/mkdocs/mkdocs.yml
@@ -0,0 +1,37 @@
+site_name: GradInf
+
+theme:
+ name: 'material'
+ home: index
+ toc_depth: 2
+ features:
+ - content.code.copy
+ - navigation.expand
+ logo: images/logo-white.png
+ palette:
+ primary: black
+
+nav:
+ - Home: index.md
+ - Examples:
+ - Differentiating a Queue Model: examples/QueueModel.md
+ - Estimating Greeks of Option Pricing Model: examples/OptionPricing.md
+ - Parameter Inference of Gene Transcription Model: examples/GeneTranscription.md
+
+markdown_extensions:
+ - pymdownx.arithmatex:
+ generic: true
+ - pymdownx.highlight:
+ anchor_linenums: true
+ line_spans: __span
+ pygments_lang_class: true
+ - pymdownx.inlinehilite
+ - pymdownx.snippets
+ - pymdownx.superfences
+
+extra_javascript:
+ - https://unpkg.com/mathjax@3/es5/tex-mml-chtml.js
+ - overrides.js
+
+extra_css:
+ - overrides.css
diff --git a/docs/requirements.txt b/docs/requirements.txt
new file mode 100644
index 0000000..f79c751
--- /dev/null
+++ b/docs/requirements.txt
@@ -0,0 +1,3 @@
+mkdocs
+mkdocs-material
+pymdown-extensions
diff --git a/tutorials/GeneTranscription.lhs b/tutorials/GeneTranscription.lhs
index 464b715..7088d86 100644
--- a/tutorials/GeneTranscription.lhs
+++ b/tutorials/GeneTranscription.lhs
@@ -18,6 +18,9 @@ and apply a loss function that measures the relative error between
the simulated mean copy numbers of mRNA and protein and ground truth
reference values.
+Writing the Model
+-----------------
+
\begin{code}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE ScopedTypeVariables #-}
@@ -33,8 +36,6 @@ import System.Random (setStdGen, mkStdGen)
import Data.List (transpose)
import Numeric.GradInf.Primitives.DeterministicPrimitives
import Numeric.GradInf.Primitives.Uniform
-import Numeric.GradInf.Primitives.Categorical
-import Numeric.GradInf.Primitives.CategoricalScore
import Numeric.GradInf.Primitives.CategoricalMaxIndep
import Numeric.GradInf.Primitives.FoldP
import Numeric.GradInf.Primitives.IterateP
@@ -54,15 +55,12 @@ n = 1000
capT :: Double
capT = 2.5
-data PrimitiveChoice = MaxIndepPrimitive | ScorePrimitive | NoAnnotationPrimitive
-
geneTranscriptionKernel ::
forall m d i b mat.
(DeterministicPrimitives d i b mat,
- Categorical m d [i], CategoricalScore m d [i],
CategoricalMaxIndep m d [i], Uniform m d)
- => PrimitiveChoice -> [d] -> ([i], d, [d]) -> m ([i], d, [d])
-geneTranscriptionKernel primitiveChoice theta (x, t, acc) = do
+ => [d] -> ([i], d, [d]) -> m ([i], d, [d])
+geneTranscriptionKernel theta (x, t, acc) = do
if extractDouble t > capT then
return (x, t, acc)
@@ -79,10 +77,7 @@ geneTranscriptionKernel primitiveChoice theta (x, t, acc) = do
let totalRate :: d = sum rates
let probs :: [d] = map (\q -> q / totalRate) rates
- x' <- case primitiveChoice of
- MaxIndepPrimitive -> categoricalMaxIndep (probs, [[m + 1, p], [m, p + 1], [m - 1, p], [m, p - 1]])
- ScorePrimitive -> categoricalScore (probs, [[m + 1, p], [m, p + 1], [m - 1, p], [m, p - 1]])
- NoAnnotationPrimitive -> categorical (probs, [[m + 1, p], [m, p + 1], [m - 1, p], [m, p - 1]])
+ x' <- categoricalMaxIndep (probs, [[m + 1, p], [m, p + 1], [m - 1, p], [m, p - 1]])
u <- uniform
let t' = t + log (1 / u) / totalRate
@@ -97,14 +92,13 @@ geneTranscriptionModel ::
forall m d i b mat.
(DeterministicPrimitives d i b mat, IterateP m ([i], d, [d]),
FoldrP m ([i], d, [d]) ([d], d),
- Categorical m d [i], CategoricalScore m d [i],
CategoricalMaxIndep m d [i], Uniform m d)
- => PrimitiveChoice -> Int -> Int -> [d] -> m d
-geneTranscriptionModel primitiveChoice numObs n theta = do
+ => Int -> Int -> [d] -> m d
+geneTranscriptionModel numObs n theta = do
let x0 = [5, 40]
let simulateTrace :: m ([i], d, [d]) =
- iterateP (geneTranscriptionKernel primitiveChoice theta) (x0, 0.0, [0.0, 0.0]) !! n
+ iterateP (geneTranscriptionKernel theta) (x0, 0.0, [0.0, 0.0]) !! n
let foldF (_, t, timeInt) (accX, accT) = do
if extractDouble t < capT then
@@ -129,10 +123,13 @@ differentiating the gene transcription model.
\begin{code}
getSamples :: Sampler [Double]
getSamples =
- replicateM 5 (geneTranscriptionModel NoAnnotationPrimitive numObs n
+ replicateM 5 (geneTranscriptionModel numObs n
(map log [18, 8, 1.5, 4]))
\end{code}
+Differentiating the Model
+-------------------------
+
We now apply GradInf to compute gradient estimates with respect
to all four parameters simultaneously using the `ZipList` functor.
This computes the full gradient in a single call to `gradInfAD`.
@@ -153,13 +150,12 @@ getGradientEstimates =
let Coupled (accPA, accPB) = acc !! 1
let horizon = fromDouble $ min 0.5 (capT - extractDouble tA)
let timeScale = 1 / 20
- let proteinAdj = 1.0
abs ((accMB - accMA) + (toDouble mB - toDouble mA) * horizon) / fromDouble 10.4
- + proteinAdj * abs ((accPB - accPA) + (toDouble pB - toDouble pA) * horizon) / fromDouble 22.3
+ + abs ((accPB - accPA) + (toDouble pB - toDouble pA) * horizon) / fromDouble 22.3
+ abs (tB - tA) * timeScale
in replicateM 10
( gradInfAD
- (geneTranscriptionModel MaxIndepPrimitive numObs n . getZipList)
+ (geneTranscriptionModel numObs n . getZipList)
( TwistedSMCInference 1 twistFunc ::
forall d i b mat. (DeterministicPrimitives d i b mat) =>
TwistedSMCInference Base CRN d i b
diff --git a/tutorials/OptionPricing.lhs b/tutorials/OptionPricing.lhs
index 284d3cc..f34bf7a 100644
--- a/tutorials/OptionPricing.lhs
+++ b/tutorials/OptionPricing.lhs
@@ -5,6 +5,9 @@ The model simulates the price of an asset in accordance with the
and returns the expected payoff of a
[European-style option](https://en.wikipedia.org/wiki/Option_style).
+Writing the Model
+-----------------
+
First, we define the model below.
\begin{code}
@@ -66,6 +69,9 @@ optionPricingModel n theta = do
0
\end{code}
+Differentiating the Model
+-------------------------
+
We now apply GradInf to compute gradient estimates.
For our inference algorithm, we will
use twisted Sequential Monte Carlo, supplying a simple twist
diff --git a/tutorials/QueueModel.lhs b/tutorials/QueueModel.lhs
index b5dec96..4de4279 100644
--- a/tutorials/QueueModel.lhs
+++ b/tutorials/QueueModel.lhs
@@ -2,7 +2,8 @@ In this example, we use GradInf to develop a new, lower-variance
gradient estimator for an
[M/M/C queue model](https://en.wikipedia.org/wiki/M/M/c_queue).
-== Setup
+Setup
+-----
This example needs the following imports.
@@ -44,7 +45,8 @@ meanAndVariance xs = do
(mu, var)
\end{code}
-== Writing the Model
+Writing the Model
+-----------------
We first write the Markov kernel for the model, which simulates
a single queuing event.
@@ -124,6 +126,9 @@ at each of these parameters.
+Differentiating the Model
+-------------------------
+
Let us now run GradInf on this program, first using an inference
scheme based on stratified importance resampling (SIR).