Fix AD (#244)

* Fix mean_vector

* Fix tests

* Bump patch

* Fix examples

Author: willtebbutt

Project.toml

@@ -1,6 +1,6 @@
 name = "Stheno"
 uuid = "8188c328-b5d6-583d-959b-9690869a5511"
-version = "0.8.1"
+version = "0.8.2"
 
 [deps]
 AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"

examples/extended_mauna_loa/script.jl

@@ -48,7 +48,7 @@ ml_df, Ttr_df, Tte_df = let
     mauna_loa_co2 = let
         mauna_loa_data_raw = CSV.read(
             joinpath(datadep"mauna_loa", "monthly_in_situ_co2_mlo.csv"), DataFrame;
-            skipto=58, header=false,
+            skipto=61, header=false,
         )
 
         data = DataFrame(

src/Stheno.jl

@@ -19,6 +19,8 @@ module Stheno
 
     const AV{T} = AbstractVector{T}
 
+    import ChainRulesCore: rrule
+
     # A couple of AbstractVector subtypes useful for expressing structure in inputs
     # regularly found in GPPPs.
     include("input_collection_types.jl")

src/affine_transformations/addition.jl

@@ -19,6 +19,10 @@ end
 
 const add_args = Tuple{typeof(+), AbstractGP, AbstractGP}
 
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(mean), ::add_args, ::AV)
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(cov), ::add_args, ::AV)
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(var), ::add_args, ::AV)
+
 mean((_, fa, fb)::add_args, x::AV) = mean(fa, x) .+ mean(fb, x)
 
 function cov((_, fa, fb)::add_args, x::AV)
@@ -62,6 +66,10 @@ end
 
 const add_known{T} = Tuple{typeof(+), T, AbstractGP}
 
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(mean), ::add_known, ::AV)
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(cov), ::add_known, ::AV)
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(var), ::add_known, ::AV)
+
 mean((_, b, f)::add_known, x::AV) = b.(x) .+ mean(f, x)
 mean((_, b, f)::add_known{<:Real}, x::AV) = b .+ mean(f, x)
 

src/affine_transformations/compose.jl

@@ -7,9 +7,12 @@ Constructs the DerivedGP f′ given by f′(x) := f(g(x))
 """
 ∘(f::AbstractGP, g) = DerivedGP((∘, f, g), f.gpc)
 
-
 const comp_args = Tuple{typeof(∘), AbstractGP, Any}
 
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(mean), ::comp_args, ::AV)
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(cov), ::comp_args, ::AV)
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(var), ::comp_args, ::AV)
+
 mean((_, f, g)::comp_args, x::AV) = mean(f, g.(x))
 
 cov((_, f, g)::comp_args, x::AV) = cov(f, g.(x))

src/affine_transformations/cross.jl

@@ -14,8 +14,8 @@ end
 _collect(X::BlockArray) = Array(X)
 
 function ChainRulesCore.rrule(::typeof(_collect), X::BlockArray)
-    function Array_pullback(Δ::Array)
-        ΔX = Tangent{Any}(blocks=BlockArray(Δ, axes(X)).blocks, axes=NoTangent())
+    function Array_pullback(Δ::AbstractArray)
+        ΔX = Tangent{Any}(blocks=BlockArray(collect(Δ), axes(X)).blocks, axes=NoTangent())
         return (NoTangent(), ΔX)
     end
     return Array(X), Array_pullback
@@ -45,9 +45,12 @@ function consistency_checks(fs)
 end
 ChainRulesCore.@non_differentiable consistency_checks(::Any)
 
-
 const cross_args{T<:AbstractVector{<:AbstractGP}} = Tuple{typeof(cross), T}
 
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(mean), ::cross_args, ::AV)
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(cov), ::cross_args, ::AV)
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(var), ::cross_args, ::AV)
+
 function mean((_, fs)::cross_args, x::BlockData)
     blks = map((f, blk)->mean(f, blk), fs, blocks(x))
     return _collect(_mortar(blks))

src/affine_transformations/product.jl

@@ -10,10 +10,14 @@ If `f isa Real`, then `h(x) = f * g(x)`.
 """
 *(f, g::AbstractGP) = DerivedGP((*, f, g), g.gpc)
 *(f::AbstractGP, g) = DerivedGP((*, g, f), f.gpc)
-*(f::AbstractGP, g::AbstractGP) = throw(ArgumentError("Cannot multiply two GPs together."))
+*(::AbstractGP, ::AbstractGP) = throw(ArgumentError("Cannot multiply two GPs together."))
 
 const prod_args{Tf} = Tuple{typeof(*), Tf, <:AbstractGP}
 
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(mean), ::prod_args, ::AV)
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(cov), ::prod_args, ::AV)
+@opt_out rrule(::RuleConfig{>:HasReverseMode}, ::typeof(var), ::prod_args, ::AV)
+
 #
 # Scale by a function
 #

src/gp/atomic_gp.jl

@@ -21,6 +21,10 @@ end
 
 atomic(gp::Tgp, gpc::GPC) where {Tgp<:AbstractGP} = AtomicGP{Tgp}(gp, gpc)
 
+@opt_out rrule(::typeof(mean), ::AtomicGP, ::AbstractVector)
+@opt_out rrule(::typeof(cov), ::AtomicGP, ::AbstractVector)
+@opt_out rrule(::typeof(var), ::AtomicGP, ::AbstractVector)
+
 AbstractGPs.mean(f::AtomicGP, x::AbstractVector) = mean(f.gp, x)
 
 AbstractGPs.cov(f::AtomicGP, x::AbstractVector) = cov(f.gp, x)

src/gp/derived_gp.jl

@@ -16,6 +16,10 @@ struct DerivedGP{Targs} <: SthenoAbstractGP
 end
 DerivedGP(args::Targs, gpc::GPC) where {Targs} = DerivedGP{Targs}(args, gpc)
 
+@opt_out rrule(::typeof(mean), ::DerivedGP, ::AbstractVector)
+@opt_out rrule(::typeof(cov), ::DerivedGP, ::AbstractVector)
+@opt_out rrule(::typeof(var), ::DerivedGP, ::AbstractVector)
+
 AbstractGPs.mean(f::DerivedGP, x::AbstractVector) = mean(f.args, x)
 
 AbstractGPs.cov(f::DerivedGP, x::AbstractVector) = cov(f.args, x)

test/gp/atomic_gp.jl

@@ -11,7 +11,7 @@ struct ToyAbstractGP <: AbstractGP end
         x = collect(range(-1.0, 1.0; length=N))
         x′ = collect(range(-1.0, 1.0; length=N′))
 
-        @test mean(f, x) == AbstractGPs._map_meanfunction(m, x)
+        @test mean(f, x) == mean_vector(m, x)
         @test cov(f, x) == kernelmatrix(k, x)
         AbstractGPs.TestUtils.test_internal_abstractgps_interface(rng, f, x, x′)
     end
@@ -27,8 +27,8 @@ struct ToyAbstractGP <: AbstractGP end
         k1, k2 = SqExponentialKernel(), SqExponentialKernel()
         f1, f2 = atomic(GP(m1, k1), gpc), atomic(GP(m2, k2), gpc)
 
-        @test mean(f1, x) == AbstractGPs._map_meanfunction(m1, x)
-        @test mean(f2, x) == AbstractGPs._map_meanfunction(m2, x)
+        @test mean(f1, x) == mean_vector(m1, x)
+        @test mean(f2, x) == mean_vector(m2, x)
 
         @test cov(f1, f2, x, x′) == zeros(N, N′)
         @test var(f1, x) == ones(N)