Merge pull request #92 from ReactiveBayes/add-autograd-backed-for-closed-form-startegy

fix: add Enzyme backend into ClosedFormStrategy

Author: Nimrais

Project.toml

@@ -28,7 +28,7 @@ ClosedFormExpectationsExt = "ClosedFormExpectations"
 [compat]
 BayesBase = "1.5.0"
 Bumper = "0.6"
-ClosedFormExpectations = "0.3.0"
+ClosedFormExpectations = "0.4.0"
 Distributions = "0.25"
 ExponentialFamily = "2.0.0"
 ExponentialFamilyManifolds = "3.0.3"
@@ -54,6 +54,7 @@ ClosedFormExpectations = "70ff922c-62d4-418d-abfc-e284e489b734"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ClosedFormExpectations = "70ff922c-62d4-418d-abfc-e284e489b734"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
@@ -66,4 +67,4 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Aqua", "BenchmarkTools", "ClosedFormExpectations", "Hwloc", "Plots", "Printf", "ForwardDiff", "Manifolds", "ReTestItems", "RollingFunctions", "JET", "StableRNGs"]
+test = ["Test", "Aqua", "BenchmarkTools", "ClosedFormExpectations", "Enzyme", "Hwloc", "Plots", "Printf", "ForwardDiff", "Manifolds", "ReTestItems", "RollingFunctions", "JET", "StableRNGs"]

docs/Project.toml

@@ -10,5 +10,6 @@ ExponentialFamilyManifolds = "5c9727c4-3b82-4ab3-b165-76e2eb971b08"
 ExponentialFamilyProjection = "17f509fa-9a96-44ba-99b2-1c5f01f0931b"
 LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"

docs/src/index.md

@@ -482,6 +482,56 @@ The `ClosedFormStrategy` typically provides:
 - **Better accuracy**: Exact gradient computations
 - **Speed advantages**: Especially significant for lower-dimensional problems
 
+### Autograd-backed closed-form strategy (Enzyme.jl)
+
+When `ClosedFormExpectations.jl` provides `ClosedFormExpectation` (i.e. ``\mathbb{E}_q[f]``) for a target–variational pair but does **not** yet have a hand-coded `ClosedWilliamsProduct` (i.e. ``\mathbb{E}_q[f \nabla_\eta \log q]``), you can still use `ClosedFormStrategy` by passing an `EnzymeBackend`. This exploits the identity
+
+```math
+\nabla_\eta \mathbb{E}_q[f(x)] = \mathbb{E}_q[f(x) \nabla_\eta \log q(x;\eta)]
+```
+
+and lets Enzyme.jl compute the gradient automatically by differentiating the closed-form expectation with respect to the natural parameters.
+
+#### Example: Gamma projected to LogNormal
+
+For this pair the `ClosedFormExpectation` is available but there is no hand-coded `ClosedWilliamsProduct`, so the default `ClosedFormStrategy()` would fail. With `EnzymeBackend` it works out of the box:
+
+```@example enzymebackend
+using ExponentialFamilyProjection, ClosedFormExpectations, Enzyme
+using ExponentialFamily, BayesBase, Distributions
+using Plots
+
+target_dist = Gamma(3.0, 2.0)
+
+result = project_to(
+    ProjectedTo(
+        LogNormal;
+        parameters = ProjectionParameters(
+            strategy = ClosedFormStrategy(EnzymeBackend()),
+            niterations = 100,
+            tolerance = 1e-6,
+        ),
+    ),
+    target_dist,
+)
+
+xs = 0.01:0.05:20.0
+
+plot(xs, x -> pdf(target_dist, x),
+     label="Target (Gamma)", linewidth=2,
+     fill=0, fillalpha=0.2, color=:blue)
+plot!(xs, x -> pdf(result, x),
+      label="Projection (LogNormal)", linewidth=2,
+      linestyle=:dash, color=:red)
+xlabel!("x")
+ylabel!("Density")
+title!("Gamma → LogNormal (EnzymeBackend)")
+```
+
+The `EnzymeBackend` supports both reverse and forward mode:
+- `ClosedFormStrategy(EnzymeBackend())` — reverse mode (default)
+- `ClosedFormStrategy(EnzymeBackend(EnzymeForward()))` — forward mode
+
 ### 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:

ext/ClosedFormExpectationsExt/ClosedFormExpectationsExt.jl

@@ -51,7 +51,7 @@ function ExponentialFamilyProjection.compute_gradient!(
     )
 
     # Compute ∇_η E[log p̃ * (T - μ)]
-    grad_target = mean(ClosedWilliamsProduct(), target_fn, q_dist)
+    grad_target = mean(ClosedWilliamsProduct(strategy.backend), target_fn, q_dist)
     grad_eta = logbasemeasure_correction(
         strategy,
         ExponentialFamily.isbasemeasureconstant(q_dist),
@@ -153,12 +153,19 @@ function ExponentialFamilyProjection.preprocess_strategy_argument(
     return (strategy, Logpdf(captured))
 end
 
-# Generic fallback for non-Function arguments
+# Wrap Distribution in Logpdf for ClosedFormStrategy
 function ExponentialFamilyProjection.preprocess_strategy_argument(
     strategy::ClosedFormStrategy,
     argument::Distribution,
 )
-    # ClosedFormStrategy accepts any callable or distribution as argument
+    return (strategy, Logpdf(argument))
+end
+
+# Wrap ProductOf in Logpdf for ClosedFormStrategy
+function ExponentialFamilyProjection.preprocess_strategy_argument(
+    strategy::ClosedFormStrategy,
+    argument::ProductOf,
+)
     return (strategy, Logpdf(argument))
 end
 
@@ -167,7 +174,6 @@ function ExponentialFamilyProjection.preprocess_strategy_argument(
     strategy::ClosedFormStrategy,
     argument,
 )
-    # ClosedFormStrategy accepts any callable or distribution as argument
     return (strategy, argument)
 end
 

src/strategies/closed_form.jl

@@ -1,14 +1,21 @@
 export ClosedFormStrategy
 
 """
-    ClosedFormStrategy <: ExponentialFamilyProjection.AbstractStrategy
+    ClosedFormStrategy{B} <: ExponentialFamilyProjection.AbstractStrategy
 
-A projection strategy that uses `ClosedFormExpectations.jl` to compute the exact gradient 
+A projection strategy that uses `ClosedFormExpectations.jl` to compute the exact gradient
 of the cross-entropy term \$\\mathbb{E}_{q_\\eta}[\\log \\tilde{p}(x)]\$ analytically.
 
-This strategy provides a "Zero-Variance" gradient estimator, avoiding the noise associated 
+This strategy provides a "Zero-Variance" gradient estimator, avoiding the noise associated
 with Monte Carlo sampling (like in `ControlVariateStrategy`).
 
+The optional `backend` field selects the differentiation backend used for computing
+`ClosedWilliamsProduct`. When `backend = nothing` (the default), hand-coded closed-form
+implementations are used. When an `EnzymeBackend` is provided, Enzyme.jl automatically
+differentiates the `ClosedFormExpectation` to obtain the Williams product gradient, enabling
+the strategy to work for any target-variational pair where the expectation is implemented
+but the Williams product is not.
+
 # Requirements
 
 To use this strategy, you **must** load the `ClosedFormExpectations` package:
@@ -17,7 +24,7 @@ To use this strategy, you **must** load the `ClosedFormExpectations` package:
 using ClosedFormExpectations
 ```
 
-Loading `ClosedFormExpectations` will trigger a package extension that implements 
+Loading `ClosedFormExpectations` will trigger a package extension that implements
 the gradient computation for this strategy.
 
 # When to Use
@@ -28,7 +35,11 @@ Use `ClosedFormStrategy` when:
 - You want faster convergence with fewer iterations
 - Reproducibility is critical (no random sampling)
 
-# Example
+Use `ClosedFormStrategy(EnzymeBackend())` when:
+- A `ClosedFormExpectation` is implemented for the pair, but `ClosedWilliamsProduct` is not
+- You want to exploit the identity \$\\nabla_\\eta \\mathbb{E}_q[f] = \\mathbb{E}_q[f \\nabla_\\eta \\log q]\$ via autodiff
+
+# Examples
 
 ```julia
 using ExponentialFamilyProjection, ClosedFormExpectations
@@ -37,7 +48,7 @@ using Distributions
 # Target distribution
 target = LogNormal(1.0, 0.5)
 
-# Project to Gamma using closed-form gradients
+# Project to Gamma using closed-form gradients (hand-coded Williams product)
 result = project_to(
     ProjectedTo(
         Gamma;
@@ -50,13 +61,37 @@ result = project_to(
 )
 ```
 
+```julia
+using ExponentialFamilyProjection, ClosedFormExpectations, Enzyme
+using Distributions
+
+# Target distribution (Gamma → LogNormal: ClosedFormExpectation is available
+# but ClosedWilliamsProduct is not, so we use EnzymeBackend to autodiff it)
+target = Gamma(2.0, 1.0)
+
+result = project_to(
+    ProjectedTo(
+        LogNormal;
+        parameters = ProjectionParameters(
+            strategy = ClosedFormStrategy(EnzymeBackend()),
+            niterations = 50
+        )
+    ),
+    Logpdf(target)
+)
+```
+
 # References
 
 This estimator was proposed in [Lukashchuk et al., 2024](https://proceedings.mlr.press/v246/lukashchuk24a.html).
 
 !!! note
-    This strategy requires that `ClosedFormExpectations.jl` implements `ClosedWilliamsProduct` 
-    for the specific pair of target distribution and variational family you're using.
-    See the `ClosedFormExpectations.jl` documentation for supported combinations.
+    Without a backend, this strategy requires that `ClosedFormExpectations.jl` implements
+    `ClosedWilliamsProduct` for the specific target-variational pair. With an `EnzymeBackend`,
+    it suffices to have `ClosedFormExpectation` implemented. See the `ClosedFormExpectations.jl`
+    documentation for supported combinations.
 """
-struct ClosedFormStrategy end
+struct ClosedFormStrategy{B}
+    backend::B
+end
+ClosedFormStrategy() = ClosedFormStrategy(nothing)