Add basic docs, expand filldist/arraydist docs

.github/workflows/CompatHelper.yml

@@ -12,5 +12,5 @@ jobs:
       - name: CompatHelper.main()
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          COMPATHELPER_PRIV: ${{ secrets.COMPATHELPER_PRIV }}
+          COMPATHELPER_PRIV: ${{ secrets.DOCUMENTER_KEY }}
         run: julia -e 'using CompatHelper; CompatHelper.main(; subdirs = ["", "test"])'

.github/workflows/Docs.yml

@@ -0,0 +1,38 @@
+name: Documentation
+
+on:
+  push:
+    branches:
+      - main
+    tags: '*'
+  pull_request:
+    branches:
+      - main
+
+concurrency:
+  # Skip intermediate builds: always.
+  # Cancel intermediate builds: only if it is a pull request build.
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
+permissions:
+  contents: write
+  pull-requests: read
+  statuses: write
+
+jobs:
+  docs:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: julia-actions/setup-julia@latest
+        with:
+          version: '1'
+      - name: Install dependencies
+        run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
+      - name: Build and deploy
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # For authentication with GitHub Actions token
+          DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} # For authentication with SSH deploy key
+          JULIA_DEBUG: Documenter # Print `@debug` statements (https://github.com/JuliaDocs/Documenter.jl/issues/955)
+        run: julia --project=docs/ docs/make.jl

.github/workflows/DocsNav.yml

@@ -0,0 +1,49 @@
+name: Add Navbar
+
+on:
+  page_build:  # Triggers the workflow on push events to gh-pages branch
+  workflow_dispatch:  # Allows manual triggering
+  schedule:
+    - cron: '0 0 * * 0'  # Runs every week on Sunday at midnight (UTC)
+
+jobs:
+  add-navbar:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - name: Checkout gh-pages
+        uses: actions/checkout@v4
+        with:
+          ref: gh-pages
+          fetch-depth: 0
+
+      - name: Download insert_navbar.sh
+        run: |
+          curl -O https://raw.githubusercontent.com/TuringLang/turinglang.github.io/main/assets/scripts/insert_navbar.sh
+          chmod +x insert_navbar.sh
+
+      - name: Update Navbar
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          git config user.name github-actions[bot]
+          git config user.email github-actions[bot]@users.noreply.github.com
+
+          # Define the URL of the navbar to be used
+          NAVBAR_URL="https://raw.githubusercontent.com/TuringLang/turinglang.github.io/main/assets/scripts/TuringNavbar.html"
+
+          # Update all HTML files in the current directory (gh-pages root)
+          ./insert_navbar.sh . $NAVBAR_URL
+
+          # Remove the insert_navbar.sh file
+          rm insert_navbar.sh
+
+          # Check if there are any changes
+          if [[ -n $(git status -s) ]]; then
+            git add .
+            git commit -m "Added navbar and removed insert_navbar.sh"
+            git push "https://${GITHUB_ACTOR}:${GITHUB_TOKEN}@github.com/${GITHUB_REPOSITORY}.git" gh-pages
+          else
+            echo "No changes to commit"
+          fi

.github/workflows/TagBot.yml

@@ -12,3 +12,4 @@ jobs:
       - uses: JuliaRegistries/TagBot@v1
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
+          ssh: ${{ secrets.DOCUMENTER_KEY }}

docs/Project.toml

@@ -0,0 +1,4 @@
+[deps]
+DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+DocumenterInterLinks = "d12716ef-a0f6-4df4-a9f1-a5a34e75c656"

docs/make.jl

@@ -0,0 +1,14 @@
+using Documenter
+using DistributionsAD
+
+makedocs(;
+    sitename="DistributionsAD",
+    modules=[DistributionsAD],
+    pages=[
+        "Home" => "index.md",
+        "API" => "api.md",
+    ],
+    doctest=false,
+)
+
+deploydocs(; repo="github.com/TuringLang/DistributionsAD.jl.git", push_preview=true)

docs/src/api.md

@@ -0,0 +1,8 @@
+# API
+
+## Functions
+
+```@docs
+filldist
+arraydist
+```

docs/src/index.md

@@ -0,0 +1,4 @@
+# DistributionsAD.jl
+
+This package defines the necessary functions to enable automatic differentiation (AD) of the `logpdf` function from [Distributions.jl](https://github.com/JuliaStats/Distributions.jl) using the packages [Tracker.jl](https://github.com/FluxML/Tracker.jl), [Zygote.jl](https://github.com/FluxML/Zygote.jl), [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) and [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl).
+The goal of this package is to make the output of `logpdf` differentiable wrt all continuous parameters of a distribution as well as the random variable in the case of continuous distributions.

src/arraydist.jl

@@ -1,7 +1,42 @@
 """
-    arraydist(dists)
+    arraydist(dists::AbstractArray{<:Distribution})
 
-Create a distribution from an array of distributions.
+Create a product distribution from an array of sub-distributions. Each element
+of `dists` should have the same size. If the size of each element is `(d1, d2,
+...)`, and `size(dists)` is `(n1, n2, ...)`, then the resulting distribution
+will have size `(d1, d2, ..., n1, n2, ...)`.
+
+The default behaviour is to directly use
+[`Distributions.product_distribution`](https://juliastats.org/Distributions.jl/stable/multivariate/#Distributions.product_distribution),
+although this can sometimes be specialised.
+
+# Examples
+
+```jldoctest; setup=:(using Distributions, Random)
+julia> d1 = arraydist([Normal(0, 1), Normal(10, 1)])
+Product{Continuous, Normal{Float64}, Vector{Normal{Float64}}}(v=Normal{Float64}[Normal{Float64}(μ=0.0, σ=1.0), Normal{Float64}(μ=10.0, σ=1.0)])
+
+julia> size(d1)
+(2,)
+
+julia> Random.seed!(42); rand(d1)
+2-element Vector{Float64}:
+ 0.7883556016042917
+ 9.1201414040456
+
+julia> d2 = arraydist([Normal(0, 1) Normal(5, 1); Normal(10, 1) Normal(15, 1)])
+DistributionsAD.MatrixOfUnivariate{Continuous, Normal{Float64}, Matrix{Normal{Float64}}}(
+dists: Normal{Float64}[Normal{Float64}(μ=0.0, σ=1.0) Normal{Float64}(μ=5.0, σ=1.0); Normal{Float64}(μ=10.0, σ=1.0) Normal{Float64}(μ=15.0, σ=1.0)]
+)
+
+julia> size(d2)
+(2, 2)
+
+julia> Random.seed!(42); rand(d2)
+2×2 Matrix{Float64}:
+ 0.788356   4.12621
+ 9.12014   14.2667
+```
 """
 arraydist(dists::AbstractArray{<:Distribution}) = product_distribution(dists)
 

src/filldist.jl

@@ -1,8 +1,36 @@
-# Default implementation just defers to Distributions.jl.
 """
     filldist(d::Distribution, ns...)
 
-Create a product distribution using `FillArrays.Fill` as the array type.
+Create a product distribution from a single distribution and a list of
+dimension sizes. If `size(d)` is `(d1, d2, ...)` and `ns` is `(n1, n2, ...)`,
+then the resulting distribution will have size `(d1, d2, ..., n1, n2, ...)`.
+
+The default behaviour is to use
+[`Distributions.product_distribution`](https://juliastats.org/Distributions.jl/stable/multivariate/#Distributions.product_distribution),
+with `FillArrays.Fill` supplied as the array argument. However, this behaviour
+is specialised in some instances, such as the one shown below.
+
+When sampling from the resulting distribution, the output will be an array where
+each element is sampled from the original distribution `d`.
+
+# Examples
+
+```jldoctest; setup=:(using Distributions, Random)
+julia> d = filldist(Normal(0, 1), 4, 5)  # size(Normal(0, 1)) == ()
+DistributionsAD.MatrixOfUnivariate{Continuous, Normal{Float64}, FillArrays.Fill{Normal{Float64}, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}(
+dists: Fill(Normal{Float64}(μ=0.0, σ=1.0), 4, 5)
+)
+
+julia> size(d)
+(4, 5)
+
+julia> Random.seed!(42); rand(d)
+4×5 Matrix{Float64}:
+ -0.363357   0.816307  -2.11433     0.433886  -0.206613
+  0.251737   0.476738   0.0437817  -0.39544   -0.310744
+ -0.314988  -0.859555  -0.825335    0.517131  -0.0404734
+ -0.311252  -1.46929    0.840289    1.44722    0.104771
+```
 """
 filldist(d::Distribution, n1::Int, ns::Int...) = product_distribution(Fill(d, n1, ns...))
 

test/Project.toml

@@ -3,6 +3,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"

test/runtests.jl

@@ -2,6 +2,7 @@ using DistributionsAD
 
 using Combinatorics
 using Distributions
+using Documenter
 using PDMats
 import LazyArrays
 
@@ -31,3 +32,16 @@ if GROUP == "All" || GROUP in ("ForwardDiff", "Zygote", "ReverseDiff", "Tracker"
     include("ad/others.jl")
     include("ad/distributions.jl")
 end
+
+# Run doctests (but not on older versions as rng seed behaves differently)
+@static if VERSION >= v"1.10"
+    @testset "doctests" begin
+        DocMeta.setdocmeta!(
+            DistributionsAD,
+            :DocTestSetup,
+            :(using DistributionsAD);
+            recursive=true,
+        )
+        doctest(DistributionsAD; manual=false)
+    end
+end