Merge branch 'main' into 63-create-integrationtestyml

Author: Nimrais

.github/workflows/CI.yml

@@ -65,7 +65,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
         with:
-          version: '1'
+          version: '1.11'
       - uses: julia-actions/cache@v2
       - name: Configure doc environment
         shell: julia --project=docs --color=yes {0}

Project.toml

@@ -30,7 +30,7 @@ ExponentialFamily = "2.0.0"
 ExponentialFamilyManifolds = "3.0.2"
 FastCholesky = "1.3"
 FillArrays = "1"
-ForwardDiff = "0.10.36"
+ForwardDiff = "0.10.36, 1"
 LinearAlgebra = "1.10"
 Manifolds = "0.11"
 ManifoldsBase = "2"

docs/Project.toml

@@ -1,10 +1,13 @@
 [deps]
 BayesBase = "b4ee3484-f114-42fe-b91c-797d54a0c67e"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ExponentialFamily = "62312e5e-252a-4322-ace9-a5f4bf9b357b"
 ExponentialFamilyManifolds = "5c9727c4-3b82-4ab3-b165-76e2eb971b08"
 ExponentialFamilyProjection = "17f509fa-9a96-44ba-99b2-1c5f01f0931b"
 LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"

docs/make.jl

@@ -16,7 +16,7 @@ makedocs(;
         canonical = "https://reactivebayes.github.io/ExponentialFamilyProjection.jl",
         edit_link = "main",
         assets = String[],
-        repolink="github.com/ReactiveBayes/ExponentialFamilyProjection.jl",
+        repolink="https://github.com/ReactiveBayes/ExponentialFamilyProjection.jl"
     ),
     pages = ["Home" => "index.md"],
 )

docs/src/index.md

@@ -162,6 +162,7 @@ using StableRNGs
 using Distributions
 using ExponentialFamily
 using ExponentialFamilyProjection
+using StatsFuns
 using Plots
 
 # 1) Generate a reproducible dataset (shared RNG)
@@ -172,8 +173,7 @@ d = input_dim + 1
 X_feat = randn(rng, n, input_dim)
 X = hcat(ones(n), X_feat)
 β_true = [0.5, 2.0, -1.5]
-σ(z) = 1 / (1 + exp(-z))
-p = map(σ, X * β_true)
+p = map(logistic, X * β_true)
 y = rand.(Ref(rng), Bernoulli.(p));
 nothing # hide
 ```
@@ -184,16 +184,7 @@ We created a binary logistic regression dataset with an intercept and fixed `rng
 # 2) Define in-place log-posterior, gradient, and Hessian
 function logpost!(out::AbstractVector{T}, β::AbstractVector{T}) where {T<:Real}
     Xβ = X * β
-    @inline function log1pexp(z)
-        z > 0 ? z + log1p(exp(-z)) : log1p(exp(z))
-    end
-    s = zero(T)
-    @inbounds for i in 1:n
-        s += y[i] * Xβ[i] - log1pexp(Xβ[i])
-    end
-    # standard normal prior on β
-    s += -0.5 * dot(β, β)
-    out[1] = s
+    out[1] = mean(y .* Xβ .- log.(1 .+ exp.(Xβ)))
     return out
 end
 
@@ -204,7 +195,8 @@ function grad!(out::AbstractVector{T}, β::AbstractVector{T}) where {T<:Real}
         pi = 1 / (1 + exp(-Xβ[i]))
         @views out[:] .+= (y[i] - pi) .* X[i, :]
     end
-    return out
+    out .= out ./ length(y)
+    return 
 end
 
 function hess!(out::AbstractMatrix{T}, β::AbstractVector{T}) where {T<:Real}
@@ -215,6 +207,7 @@ function hess!(out::AbstractMatrix{T}, β::AbstractVector{T}) where {T<:Real}
         wi = pi * (1 - pi)
         @views out .-= wi .* (X[i, :] * transpose(X[i, :]))
     end
+    out .= out ./ length(y)
     return out
 end
 ```
@@ -358,6 +351,67 @@ plt_mc
 plot(plt_mean, plt_mc; layout = (1, 2), size = (1100, 450))
 ```
 
+### How to use autograd with Gauss–Newton (Enzyme.jl)
+
+You do not need to hand-derive gradients or Hessians to use Gauss–Newton. With `Enzyme.jl`, you can automatically obtain both and use them through the same in-place API shown above. In practice, this is typically faster and yields more stable estimates than naïve manual derivatives. `Enzyme.jl` has some sharp edges; please consult the [Enzyme documentation](https://enzymejs.github.io/enzyme/) before use.
+
+```@example gaussnewton
+using Enzyme
+using BenchmarkTools
+
+# 10) Define the log-posterior for logistic regression with a standard normal prior
+function obj(β::AbstractVector, X::AbstractMatrix, y::AbstractVector)
+    Xβ = X * β
+    return mean(y .* Xβ .- log.(1 .+ exp.(Xβ)))
+end
+
+# Reverse-mode gradient and forward-over-reverse Hessian via Enzyme
+grad_enzyme = (β, X, y) -> Enzyme.gradient(Reverse, obj, β, Const(X), Const(y))[1]
+function jacobian_enzyme(β, X, y)
+    Enzyme.jacobian(set_runtime_activity(Forward), grad_enzyme, β, Const(X), Const(y))
+end
+
+# 11) In-place wrappers expected by Gauss–Newton
+function make_logpost!(X, y)
+    (out, β) -> (out[1] = obj(β, X, y); out)
+end
+function make_grad!(X, y)
+    function _grad!(out::AbstractVector{T}, β::AbstractVector{T}) where {T}
+        out .= grad_enzyme(β, X, y)
+        return out
+    end
+    _grad!
+end
+function make_hess!(X, y)
+    function _hess!(out::AbstractMatrix{T}, β::AbstractVector{T}) where {T}
+        J, _ = jacobian_enzyme(β, X, y)
+        out .= J
+        return out
+    end
+    _hess!
+end
+
+logpostE! = make_logpost!(X, y)
+gradE! = make_grad!(X, y)
+hessE! = make_hess!(X, y)
+
+inplace_enzyme = ExponentialFamilyProjection.InplaceLogpdfGradHess(logpostE!, gradE!, hessE!)
+prj_enzyme = ProjectedTo(MvNormalMeanCovariance, d; parameters = params)
+result_enzyme = project_to(prj_enzyme, inplace_enzyme)
+```
+
+We can quickly compare the runtime of the Enzyme-based implementation to the manual one defined above.
+
+```@example gaussnewton
+# 12) Speed comparison against the manual implementation from above
+t_manual = @belapsed project_to($prj, $inplace)
+t_enzyme = @belapsed project_to($prj_enzyme, $inplace_enzyme)
+speedup = t_manual / t_enzyme
+round.((speedup, t_manual, t_enzyme); digits = 3)
+```
+
+On typical runs we observe a substantial speedup (often around 10×) for Enzyme while maintaining the same result.
+
 ### Projection with samples
 
 The projection can be done given a set of samples instead of the function directly. For example, let's project an set of samples onto a Beta distribution:

src/projected_to.jl

@@ -42,7 +42,7 @@ struct ProjectedTo{T,D,C,P,E}
 end
 
 ProjectedTo(
-    dims::Vararg{Int};
+    dims::Tuple{Vararg{Int}};
     conditioner = nothing,
     parameters = DefaultProjectionParameters(),
     kwargs = nothing,
@@ -53,6 +53,29 @@ ProjectedTo(
     parameters = parameters,
     kwargs = kwargs,
 )
+ProjectedTo(;
+    conditioner = nothing,
+    parameters = DefaultProjectionParameters(),
+    kwargs = nothing,
+) = ProjectedTo(
+    ExponentialFamilyDistribution,
+    ()...,
+    conditioner = conditioner,
+    parameters = parameters,
+    kwargs = kwargs,
+)
+ProjectedTo(
+    dim::Int;
+    conditioner = nothing,
+    parameters = DefaultProjectionParameters(),
+    kwargs = nothing,
+) = ProjectedTo(
+    ExponentialFamilyDistribution,
+    dim,
+    conditioner = conditioner,
+    parameters = parameters,
+    kwargs = kwargs,
+)
 function ProjectedTo(
     ::Type{T},
     dims...;

src/strategies/bonnet/bonnet_logpdf.jl

@@ -35,9 +35,9 @@ implementation and returns an `InplaceLogpdfGradHess` instance.
 # See also
 - `NaiveGradHess` — adapter that combines separate `grad!`/`hess!` into `grad_hess!`.
 """
-function InplaceLogpdfGradHess(logpdf!::F, grad!::G, hess!::H) where {F,G,H}
+function InplaceLogpdfGradHess(__logpdf::F, grad!::G, hess!::H) where {F,G,H}
     wrapper_grad_hess! = NaiveGradHess(grad!, hess!)
-    return InplaceLogpdfGradHess(logpdf!, wrapper_grad_hess!)
+    return InplaceLogpdfGradHess(__logpdf, wrapper_grad_hess!)
 end
 
 """