reactivate AD tests: mean functions

src/mean_function.jl

@@ -1,5 +1,7 @@
 abstract type MeanFunction end
 
+# (m::MeanFunction)(x::AbstractVector) = _map_meanfunction(m, x)
+
 """
     ZeroMean{T<:Real} <: MeanFunction
 
@@ -10,7 +12,7 @@ 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())
@@ -28,7 +30,7 @@ struct ConstMean{T<:Real} <: MeanFunction
     c::T
 end
 
-_map_meanfunction(m::ConstMean, x::AbstractVector) = Fill(m.c, length(x))
+_map_meanfunction(m::ConstMean, x::AbstractVector) = fill(m.c, length(x))
 
 """
     CustomMean{Tf} <: MeanFunction

test/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"

test/mean_function.jl

@@ -1,47 +1,52 @@
 @testset "mean_functions" begin
     @testset "ZeroMean" begin
-        P = 3
-        Q = 2
-        D = 4
-        # X = ColVecs(randn(rng, D, P))
-        x = randn(P)
-        x̄ = randn(P)
+        rng, D, N = MersenneTwister(123456), 5, 3
+        # X = ColVecs(randn(rng, D, N))
+        x = randn(rng, N)
+        x̄ = randn(rng, N)
         f = ZeroMean{Float64}()
 
         for x in [x]
             @test AbstractGPs._map_meanfunction(f, x) == zeros(size(x))
-            # differentiable_mean_function_tests(f, randn(rng, P), x)
+            differentiable_mean_function_tests(f, randn(rng, N), x)
         end
 
         # 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, f, x)
         @test y == AbstractGPs._map_meanfunction(f, x)
-        Δmap, Δf, Δx = pb(randn(P))
+        Δmap, Δf, Δx = pb(randn(rng, N))
         @test iszero(Δmap)
         @test iszero(Δf)
         @test iszero(Δx)
     end
     @testset "ConstMean" begin
-        rng, D, N = MersenneTwister(123456), 5, 3
-        # X = ColVecs(randn(rng, D, N))
-        x = randn(rng, N)
+        rng, N, D = MersenneTwister(123456), 5, 3
+        x1 = randn(rng, N)
+        xD = ColVecs(randn(rng, D, N))
+        xD′ = RowVecs(randn(rng, N, D))
+
         c = randn(rng)
         m = ConstMean(c)
 
-        for x in [x]
+        for x in [x1, xD, xD′]
             @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
     @testset "CustomMean" begin
-        rng, N, D = MersenneTwister(123456), 11, 2
-        x = randn(rng, N)
+        rng, N, D = MersenneTwister(123456), 5, 3
+        x1 = randn(rng, N)
+        xD = ColVecs(randn(rng, D, N))
+        xD′ = RowVecs(randn(rng, N, D))
+
         foo_mean = x -> sum(abs2, x)
-        f = CustomMean(foo_mean)
+        m = CustomMean(foo_mean)
 
-        @test AbstractGPs._map_meanfunction(f, x) == map(foo_mean, x)
-        # differentiable_mean_function_tests(f, randn(rng, N), x)
+        for x in [x1, xD, xD′]
+            @test AbstractGPs._map_meanfunction(m, x) == map(foo_mean, 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 -> 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
 )