reactivate AD tests: mean functions

Project.toml

@@ -1,7 +1,7 @@
 name = "AbstractGPs"
 uuid = "99985d1d-32ba-4be9-9821-2ec096f28918"
 authors = ["JuliaGaussianProcesses Team"]
-version = "0.5.11"
+version = "0.5.12"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/mean_function.jl

@@ -10,10 +10,10 @@ struct ZeroMean{T<:Real} <: MeanFunction end
 """
 This is an AbstractGPs-internal workaround for AD issues; ideally we would just extend Base.map
 """
-_map_meanfunction(::ZeroMean{T}, x::AbstractVector) where {T} = Zeros{T}(length(x))
+ _map_meanfunction(::ZeroMean{T}, x::AbstractVector) where {T} = Zeros{T}(length(x))
 
 function ChainRulesCore.rrule(::typeof(_map_meanfunction), m::ZeroMean, x::AbstractVector)
-    map_ZeroMean_pullback(Δ) = (NoTangent(), NoTangent(), ZeroTangent())
+    map_ZeroMean_pullback(Δ) = (NoTangent(), ZeroTangent(), ZeroTangent())
     return _map_meanfunction(m, x), map_ZeroMean_pullback
 end
 
@@ -30,6 +30,11 @@ end
 
 _map_meanfunction(m::ConstMean, x::AbstractVector) = Fill(m.c, length(x))
 
+function ChainRulesCore.rrule(::typeof(_map_meanfunction), m::ConstMean, x::AbstractVector)
+    map_ConstMean_pullback(Δ) = (NoTangent(), Tangent{ConstMean}(; c=sum(Δ)), ZeroTangent())
+    return _map_meanfunction(m, x), map_ConstMean_pullback
+end
+
 """
     CustomMean{Tf} <: MeanFunction
 
@@ -40,4 +45,4 @@ struct CustomMean{Tf} <: MeanFunction
     f::Tf
 end
 
-_map_meanfunction(f::CustomMean, x::AbstractVector) = map(f.f, x)
+_map_meanfunction(m::CustomMean, x::AbstractVector) = map(m.f, x)

test/mean_function.jl

@@ -10,17 +10,7 @@
 
         for x in [x1, xD_colvecs, xD_rowvecs]
             @test AbstractGPs._map_meanfunction(m, x) == zeros(N)
-            #differentiable_mean_function_tests(m, randn(rng, N), x)
-
-            # Manually verify the ChainRule. Really, this should employ FiniteDifferences, but
-            # currently ChainRulesTestUtils isn't up to handling this, so this will have to do
-            # for now.
-            y, pb = rrule(AbstractGPs._map_meanfunction, m, x)
-            @test y == AbstractGPs._map_meanfunction(m, x)
-            Δmap, Δf, Δx = pb(randn(rng, N))
-            @test iszero(Δmap)
-            @test iszero(Δf)
-            @test iszero(Δx)
+            differentiable_mean_function_tests(m, randn(rng, N), x)
         end
     end
 
@@ -30,7 +20,7 @@
 
         for x in [x1, xD_colvecs, xD_rowvecs]
             @test AbstractGPs._map_meanfunction(m, x) == fill(c, N)
-            #differentiable_mean_function_tests(m, randn(rng, N), x)
+            differentiable_mean_function_tests(m, randn(rng, N), x)
         end
     end
 
@@ -40,7 +30,7 @@
 
         for x in [x1, xD_colvecs, xD_rowvecs]
             @test AbstractGPs._map_meanfunction(m, x) == map(foo_mean, x)
-            #differentiable_mean_function_tests(m, randn(rng, N), x)
+            differentiable_mean_function_tests(m, randn(rng, N), x)
         end
     end
 end

test/test_util.jl

@@ -73,8 +73,9 @@ end
 Test _very_ basic consistency properties of the mean function `m`.
 """
 function mean_function_tests(m::MeanFunction, x::AbstractVector)
-    @test AbstractGPs._map_meanfunction(m, x) isa AbstractVector
-    @test length(ew(m, x)) == length(x)
+    mean = AbstractGPs._map_meanfunction(m, x)
+    @test mean isa AbstractVector
+    @test length(mean) == length(x)
 end
 
 """
@@ -87,34 +88,17 @@ end
 Ensure that the gradient w.r.t. the inputs of `MeanFunction` `m` are approximately correct.
 """
 function differentiable_mean_function_tests(
-    m::MeanFunction,
-    ȳ::AbstractVector{<:Real},
-    x::AbstractVector{<:Real};
-    rtol=_rtol,
-    atol=_atol,
+    m::MeanFunction, ȳ::AbstractVector, x::AbstractVector; rtol=_rtol, atol=_atol
 )
     # Run forward tests.
     mean_function_tests(m, x)
 
     # Check adjoint.
     @assert length(ȳ) == length(x)
-    return adjoint_test(x -> ew(m, x), ȳ, x; rtol=rtol, atol=atol)
+    adjoint_test(x -> collect(AbstractGPs._map_meanfunction(m, x)), ȳ, x; rtol=rtol, atol=atol)
+    return nothing
 end
 
-# function differentiable_mean_function_tests(
-#     m::MeanFunction,
-#     ȳ::AbstractVector{<:Real},
-#     x::ColVecs{<:Real};
-#     rtol=_rtol,
-#     atol=_atol,
-# )
-#     # Run forward tests.
-#     mean_function_tests(m, x)
-
-#     @assert length(ȳ) == length(x)
-#     adjoint_test(X->ew(m, ColVecs(X)), ȳ, x.X; rtol=rtol, atol=atol)  
-# end
-
 function differentiable_mean_function_tests(
     rng::AbstractRNG, m::MeanFunction, x::AbstractVector; rtol=_rtol, atol=_atol
 )