Tweak Examples to use DI (#139)

* Tweak examples

* Capitalisation

* Update examples CI

Author: willtebbutt

.github/workflows/examples.yml

@@ -32,7 +32,7 @@ jobs:
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
-      - uses: actions/cache@v1
+      - uses: actions/cache@v3
         env:
           cache-name: cache-artifacts
         with:

examples/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 KernelFunctions = "ec8451be-7e33-11e9-00cf-bbf324bd1392"

examples/approx_space_time_learning.jl

@@ -11,9 +11,10 @@ using TemporalGPs
 using TemporalGPs: Separable, approx_posterior_marginals, RegularInTime
 
 # Load standard packages from the Julia ecosystem
+using ADTypes # Way to specify algorithmic differentiation backend.
 using Optim # Standard optimisation algorithms.
 using ParameterHandling # Helper functionality for dealing with model parameters.
-using Mooncake # Algorithmic Differentiation
+import Mooncake # Algorithmic differentiation.
 
 using ParameterHandling: flatten
 
@@ -60,24 +61,19 @@ function objective(flat_params)
     return -elbo(f(x, params.var_noise), y, z_r)
 end
 
-# Optimise using Optim.
-function objective_grad(rule, flat_params)
-    return Mooncake.value_and_gradient!!(rule, objective, flat_params)[2][2]
-end
-
 @info "running objective"
 @show objective(flat_initial_params)
 
 training_results = Optim.optimize(
     objective,
-    Base.Fix1(objective_grad, Mooncake.build_rrule(objective, flat_initial_params)),
     flat_initial_params + randn(4), # Add some noise to make learning non-trivial
     BFGS(
         alphaguess = Optim.LineSearches.InitialStatic(scaled=true),
         linesearch = Optim.LineSearches.BackTracking(),
     ),
     Optim.Options(show_trace = true);
     inplace=false,
+    autodiff=AutoMooncake(; config=nothing),
 );
 
 

examples/exact_space_time_learning.jl

@@ -11,9 +11,10 @@ using TemporalGPs
 using TemporalGPs: Separable, RectilinearGrid
 
 # Load standard packages from the Julia ecosystem
+using ADTypes # Way to specify algorithmic differentiation backend.
 using Optim # Standard optimisation algorithms.
 using ParameterHandling # Helper functionality for dealing with model parameters.
-using Mooncake # Algorithmic Differentiation
+import Mooncake # Algorithmic differentiation.
 
 # Declare model parameters using `ParameterHandling.jl` types.
 flat_initial_params, unflatten = ParameterHandling.flatten((
@@ -53,21 +54,17 @@ function objective(flat_params)
     return -logpdf(f(x, params.var_noise), y)
 end
 
-function objective_grad(rule, flat_params)
-    return Mooncake.value_and_gradient!!(rule, objective, flat_params)[2][2]
-end
-
 # Optimise using Optim.
 training_results = Optim.optimize(
     objective,
-    Base.Fix1(objective_grad, Mooncake.build_rrule(objective, flat_initial_params)),
     flat_initial_params + randn(4), # Add some noise to make learning non-trivial
     BFGS(
         alphaguess = Optim.LineSearches.InitialStatic(scaled=true),
         linesearch = Optim.LineSearches.BackTracking(),
     ),
     Optim.Options(show_trace = true);
     inplace=false,
+    autodiff=AutoMooncake(; config=nothing),
 );
 
 # Extracting the final values of the parameters.

examples/exact_time_learning.jl

@@ -7,9 +7,10 @@ using AbstractGPs
 using TemporalGPs
 
 # Load standard packages from the Julia ecosystem
+using ADTypes # Way to specify algorithmic differentiation backend.
 using Optim # Standard optimisation algorithms.
 using ParameterHandling # Helper functionality for dealing with model parameters.
-using Mooncake # Algorithmic Differentiation
+import Mooncake # Algorithmic differentiation.
 
 # Declare model parameters using `ParameterHandling.jl` types.
 # var_kernel is the variance of the kernel, λ the inverse length scale, and var_noise the
@@ -48,22 +49,17 @@ function objective(flat_params)
     return -logpdf(f(x, params.var_noise), y)
 end
 
-# A helper function to get the gradient.
-function objective_grad(rule, flat_params)
-    return Mooncake.value_and_gradient!!(rule, objective, flat_params)[2][2]
-end
-
 # Optimise using Optim. Mooncake takes a little while to compile.
 training_results = Optim.optimize(
     objective,
-    Base.Fix1(objective_grad, Mooncake.build_rrule(objective, flat_initial_params)),
     flat_initial_params .+ randn.(), # Perturb the parameters to make learning non-trivial
     BFGS(
         alphaguess = Optim.LineSearches.InitialStatic(scaled=true),
         linesearch = Optim.LineSearches.BackTracking(),
     ),
     Optim.Options(show_trace = true);
     inplace=false,
+    autodiff=AutoMooncake(; config=nothing),
 );
 
 # Extracting the final values of the parameters. Should be moderately close to truth.