LogDensityProblemsAD + fix CI bug (#123)

* Make temporalgps run in CI

* LogDensityProblemsAD interface

* Update test/runtests.jl

* Remove ADTypes test for now

* Sort out ordering in Project.toml

* Bump patch

---------

Co-authored-by: Hong Ge <[email protected]>

.github/workflows/CI.yml

@@ -30,6 +30,7 @@ jobs:
           - 'integration_testing/special_functions'
           - 'integration_testing/array'
           - 'integration_testing/turing'
+          - 'integration_testing/temporalgps'
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v1

Project.toml

@@ -1,7 +1,7 @@
 name = "Tapir"
 uuid = "07d77754-e150-4737-8c94-cd238a1fb45b"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -17,9 +17,11 @@ Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [weakdeps]
+LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [extensions]
+TapirLogDensityProblemsADExt = "LogDensityProblemsAD"
 TapirSpecialFunctionsExt = "SpecialFunctions"
 
 [compat]
@@ -32,6 +34,7 @@ ExprTools = "0.1"
 FillArrays = "1"
 Graphs = "1"
 JET = "0.8"
+LogDensityProblemsAD = "1"
 PDMats = "0.11"
 Setfield = "1"
 SpecialFunctions = "2"
@@ -47,6 +50,7 @@ DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 KernelFunctions = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
+LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
@@ -56,4 +60,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [targets]
-test = ["AbstractGPs", "BenchmarkTools", "DiffTests", "Distributions", "FillArrays", "KernelFunctions", "PDMats", "ReverseDiff", "SpecialFunctions", "StableRNGs", "Test", "Turing", "TemporalGPs"]
+test = ["AbstractGPs", "BenchmarkTools", "DiffTests", "Distributions", "FillArrays", "KernelFunctions", "LogDensityProblemsAD", "PDMats", "ReverseDiff", "SpecialFunctions", "StableRNGs", "Test", "Turing", "TemporalGPs"]

ext/TapirLogDensityProblemsADExt.jl

@@ -0,0 +1,55 @@
+# This file is largely copy + pasted + modified from the Zygote extension in
+# LogDensityProblemsAD.jl. 
+
+module TapirLogDensityProblemsADExt
+
+if isdefined(Base, :get_extension)
+    using LogDensityProblemsAD: ADGradientWrapper
+    import LogDensityProblemsAD: ADgradient, logdensity_and_gradient, dimension, logdensity
+    import Tapir
+else
+    using ..LogDensityProblemsAD: ADGradientWrapper
+    import ..LogDensityProblemsAD: ADgradient, logdensity_and_gradient, dimension, logdensity
+    import ..Tapir
+end
+
+struct TapirGradientLogDensity{Trule, L} <: ADGradientWrapper
+    rule::Trule
+    l::L
+end
+
+dimension(∇l::TapirGradientLogDensity) = dimension(Tapir.primal(∇l.l))
+
+function logdensity(∇l::TapirGradientLogDensity, x::Vector{Float64})
+    return logdensity(Tapir.primal(∇l.l), x)
+end
+
+"""
+    ADgradient(Val(:Tapir), ℓ)
+
+Gradient using algorithmic/automatic differentiation via Tapir.
+"""
+function ADgradient(::Val{:Tapir}, l)
+    primal_sig = Tuple{typeof(logdensity), typeof(l), Vector{Float64}}
+    rule = Tapir.build_rrule(Tapir.TapirInterpreter(), primal_sig)
+    return TapirGradientLogDensity(rule, Tapir.uninit_fcodual(l))
+end
+
+Base.show(io::IO, ∇ℓ::TapirGradientLogDensity) = print(io, "Tapir AD wrapper for ", ∇ℓ.ℓ)
+
+# We only test Tapir with `Float64`s at the minute, so make strong assumptions about the
+# types supported in order to prevent silent errors.
+function logdensity_and_gradient(::TapirGradientLogDensity, ::AbstractVector)
+    msg = "Only Vector{Float64} presently supported for logdensity_and_gradients."
+    throw(ArgumentError(msg))
+end
+
+function logdensity_and_gradient(∇l::TapirGradientLogDensity, x::Vector{Float64})
+    dx = zeros(length(x))
+    y, pb!! = ∇l.rule(Tapir.zero_fcodual(logdensity), ∇l.l, Tapir.CoDual(x, dx))
+    @assert Tapir.primal(y) isa Float64
+    pb!!(1.0)
+    return Tapir.primal(y), dx
+end
+
+end

test/integration_testing/logdensityproblemsad_interop.jl

@@ -0,0 +1,22 @@
+using LogDensityProblemsAD
+using LogDensityProblemsAD: logdensity_and_gradient, capabilities, dimension, logdensity
+
+# Copied over from LogDensityProblemsAD test suite.
+struct TestLogDensity2 end
+LogDensityProblemsAD.logdensity(::TestLogDensity2, x) = -sum(abs2, x)
+LogDensityProblemsAD.dimension(::TestLogDensity2) = 20
+test_gradient(x) = -2 .* x
+
+@testset "AD via Zygote" begin
+    l = TestLogDensity2()
+    ∇l = ADgradient(Val(:Tapir), l)
+
+    @test dimension(∇l) == 20
+    @test capabilities(∇l) == LogDensityProblemsAD.LogDensityOrder(1)
+    for _ in 1:100
+        x = randn(20)
+        @test isapprox(@inferred(logdensity(∇l, x)), logdensity(l, x))
+        @test isapprox(logdensity_and_gradient(∇l, x)[1], logdensity(TestLogDensity2(), x))
+        @test isapprox(logdensity_and_gradient(∇l, x)[2], test_gradient(x))
+    end
+end

test/runtests.jl

@@ -41,6 +41,7 @@ include("front_matter.jl")
         end
         include("chain_rules_macro.jl")
     elseif test_group == "integration_testing/misc"
+        include(joinpath("integration_testing", "logdensityproblemsad_interop.jl"))
         include(joinpath("integration_testing/", "misc.jl"))
         include(joinpath("integration_testing", "battery_tests.jl"))
     elseif test_group == "integration_testing/diff_tests"
@@ -55,6 +56,8 @@ include("front_matter.jl")
         include(joinpath("integration_testing", "array.jl"))
     elseif test_group == "integration_testing/turing"
         include(joinpath("integration_testing", "turing.jl"))
+    elseif test_group == "integration_testing/temporalgps"
+        include(joinpath("integration_testing", "temporalgps.jl"))
     else
         throw(error("test_group=$(test_group) is not recognised"))
     end