Resolve issues with Zygote.jl and mean_vector (#415)

* resolve issues with AD of mean_vector

* add tests for mean_vector with ColVecs & RowVecs

* bump version to 0.5.22

Author: soldasim

Project.toml

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

src/mean_function.jl

@@ -44,23 +44,12 @@ mean_vector(m::ConstMean, x::AbstractVector) = Fill(m.c, length(x))
 
 A wrapper around whatever unary function you fancy. Must be able to be mapped over an
 `AbstractVector` of inputs.
-
-# Warning
-`CustomMean` is generally sufficient for testing purposes, but care should be taken if
-attempting to differentiate through `mean_vector` with a `CustomMean` when using
-`Zygote.jl`. In particular, `mean_vector(m::CustomMean, x)` is implemented as `map(m.f, x)`,
-which when `x` is a `ColVecs` or `RowVecs` will not differentiate correctly.
-
-In such cases, you should implement `mean_vector` directly for your custom mean.
-For example, if `f(x) = sum(x)`, you might implement `mean_vector` as
-```julia
-mean_vector(::CustomMean{typeof(f)}, x::ColVecs) = vec(sum(x.X; dims=1))
-mean_vector(::CustomMean{typeof(f)}, x::RowVecs) = vec(sum(x.X; dims=2))
-```
-which avoids ever applying `map` to a `ColVecs` or `RowVecs`.
 """
 struct CustomMean{Tf} <: MeanFunction
     f::Tf
 end
 
 mean_vector(m::CustomMean, x::AbstractVector) = map(m.f, x)
+
+mean_vector(m::CustomMean, x::ColVecs) = map(m.f, eachcol(x.X))
+mean_vector(m::CustomMean, x::RowVecs) = map(m.f, eachrow(x.X))

test/mean_function.jl

@@ -24,4 +24,34 @@
             differentiable_mean_function_tests(rng, m, x)
         end
     end
+
+    @testset "ColVecs & RowVecs" begin
+        m = custom_mean_testcase.mean_function
+
+        @test mean_vector(m, xD_colvecs) == map(foo_mean, eachcol(xD_colvecs.X))
+        @test mean_vector(m, xD_rowvecs) == map(foo_mean, eachrow(xD_rowvecs.X))
+    end
+
+    # This test fails without the specialized methods
+    #   `mean_vector(m::CustomMean, x::ColVecs)`
+    #   `mean_vector(m::CustomMean, x::RowVecs)`
+    @testset "Zygote gradients" begin
+        X = [1.;; 2.;; 3.;;]
+        y = [1., 2., 3.]
+        foo_mean = x -> sum(abs2, x)
+
+        function construct_finite_gp(X, lengthscale, noise)
+            mean = CustomMean(foo_mean)
+            kernel = with_lengthscale(Matern52Kernel(), lengthscale)
+            return GP(mean, kernel)(X, noise)
+        end
+
+        function loglike(lengthscale, noise)
+            gp = construct_finite_gp(X, lengthscale, noise)
+            return logpdf(gp, y)
+        end
+
+        @test Zygote.gradient(n -> loglike(1., n), 1.)[1] isa Real
+        @test Zygote.gradient(l -> loglike(l, 1.), 1.)[1] isa Real    
+    end
 end