# Zygote.gradient does not work with AbstractGPs.CustomMean

**URL:** <https://discourse.julialang.org/t/zygote-gradient-does-not-work-with-abstractgps-custommean/87815>\
**Category:** General Usage\
**Tags:** zygote, gaussian-process, gradient\
**Created:** [September 26, 2022, 2:07pm UTC](https://discourse.julialang.org/t/zygote-gradient-does-not-work-with-abstractgps-custommean/87815 "2022-09-26T14:07:30Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![soldasim](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/soldasim/32/42514_2.png) [@soldasim](https://discourse.julialang.org/u/soldasim)\
**Post date:** [September 26, 2022, 2:07pm UTC](https://discourse.julialang.org/t/zygote-gradient-does-not-work-with-abstractgps-custommean/87815/1 "2022-09-26T14:07:30Z")

</div>

Hello, I am trying to use an `AbstractGPs.FiniteGP` with Turing and I am having problems with `Zygote.gradient` calculation of the GP hyperparameters.

For some reason, Zygote fails to calculate the gradient when using an `AbstractGPs.CustomMean` but works properly when using `AbstractGPs.ZeroMean`.

When I try using ForwardDiff instead, both gradients are calculated without an issue.

MWE of the issue:

```julia
using AbstractGPs
using Zygote
using ForwardDiff

X = [1.; 2.; 3.;;]
y = [1., 2., 3.]

function construct_finite_gp(X, lengthscale, noise; min_param_val=1e-6, mean)
    # for numerical stability
    lengthscale = lengthscale + min_param_val
    noise = noise + min_param_val

    kernel = with_lengthscale(Matern52Kernel(), lengthscale)
    return GP(mean, kernel)(X', noise)
end

function f(noise; mean)
    gp = construct_finite_gp(X, 1., first(noise); mean)
    return logpdf(gp, y)
end

forwarddiff_zeromean = ForwardDiff.gradient(n -> f(n; mean=AbstractGPs.ZeroMean()), [1.])
@show forwarddiff_zeromean

forwarddiff_custommean = ForwardDiff.gradient(n -> f(n; mean=AbstractGPs.CustomMean(x->0.)), [1.])
@show forwarddiff_custommean

zygote_zeromean = Zygote.gradient(n -> f(n; mean=AbstractGPs.ZeroMean()), 1.)
@show zygote_zeromean

zygote_custommean = Zygote.gradient(n -> f(n; mean=AbstractGPs.CustomMean(x->0.)), 1.)
@show zygote_custommean

```

The result of running the MWE:

```julia
forwarddiff_zeromean = [0.2630913490608584]
forwarddiff_custommean = [0.2630913490608584]
zygote_zeromean = (0.26309134906085807,)
ERROR: MethodError: no method matching +(::NamedTuple{(:X,), Tuple{Matrix{Float64}}}, ::Vector{Nothing})
Closest candidates are:
  +(::Any, ::Any, ::Any, ::Any...) at operators.jl:591
  +(::Distributions.MvNormal, ::AbstractVector) at C:\Users\sheld\.julia\packages\Distributions\7iOJp\src\multivariate\mvnormal.jl:294
  +(::SparseArrays.AbstractSparseMatrixCSC, ::Array) at C:\Users\sheld\AppData\Local\Programs\Julia-1.8.0\share\julia\stdlib\v1.8\SparseArrays\src\sparsematrix.jl:1832
  ...
Stacktrace:
  [1] accum(x::NamedTuple{(:X,), Tuple{Matrix{Float64}}}, y::Vector{Nothing}) (repeats 2 times)
    @ Zygote C:\Users\sheld\.julia\packages\Zygote\dABKa\src\lib\lib.jl:17
  [2] (::typeof(∂(mean_and_cov)))(Δ::Tuple{Vector{Float64}, LinearAlgebra.UpperTriangular{Float64, Matrix{Float64}}}) @ Zygote C:\Users\sheld\.julia\packages\Zygote\dABKa\src\compiler\interface2.jl:0
  [3] Pullback
    @ C:\Users\sheld\.julia\packages\AbstractGPs\e2ey8\src\finite_gp_projection.jl:134 [inlined]
  [4] (::typeof(∂(mean_and_cov)))(Δ::Tuple{Vector{Float64}, LinearAlgebra.UpperTriangular{Float64, Matrix{Float64}}}) @ Zygote C:\Users\sheld\.julia\packages\Zygote\dABKa\src\compiler\interface2.jl:0
  [5] Pullback
    @ C:\Users\sheld\.julia\packages\AbstractGPs\e2ey8\src\finite_gp_projection.jl:307 [inlined]
  [6] (::typeof(∂(logpdf)))(Δ::Float64)
    @ Zygote C:\Users\sheld\.julia\packages\Zygote\dABKa\src\compiler\interface2.jl:0
  [7] Pullback
    @ d:\plzen\sandbox\sampling_error_mwe\zygote_mwe.jl:19 [inlined]
  [8] (::typeof(∂(#f#47)))(Δ::Float64)
    @ Zygote C:\Users\sheld\.julia\packages\Zygote\dABKa\src\compiler\interface2.jl:0
  [9] Pullback
    @ d:\plzen\sandbox\sampling_error_mwe\zygote_mwe.jl:17 [inlined]
 [10] (::typeof(∂(f##kw)))(Δ::Float64)
    @ Zygote C:\Users\sheld\.julia\packages\Zygote\dABKa\src\compiler\interface2.jl:0
 [11] Pullback
    @ d:\plzen\sandbox\sampling_error_mwe\zygote_mwe.jl:31 [inlined]
 [12] (::typeof(∂(#56)))(Δ::Float64)
    @ Zygote C:\Users\sheld\.julia\packages\Zygote\dABKa\src\compiler\interface2.jl:0
 [13] (::Zygote.var"#60#61"{typeof(∂(#56))})(Δ::Float64)
    @ Zygote C:\Users\sheld\.julia\packages\Zygote\dABKa\src\compiler\interface.jl:45
 [14] gradient(f::Function, args::Float64)
    @ Zygote C:\Users\sheld\.julia\packages\Zygote\dABKa\src\compiler\interface.jl:97
 [15] top-level scope
    @ d:\plzen\sandbox\sampling_error_mwe\zygote_mwe.jl:31

```

I need to use Zygote because of [this issue](https://discourse.julialang.org/t/gaussian-process-regression-with-turing-gets-stuck/86892).

Does anyone know what might be the cause of this?  
Is it a bug in Zygote or AbstractGPs or some error on my side?

---

<div class="post-metadata">

**Author:** ![willtebbutt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/willtebbutt/32/6790_2.png) [@willtebbutt](https://discourse.julialang.org/u/willtebbutt)\
**Post date:** [September 26, 2022, 2:34pm UTC](https://discourse.julialang.org/t/zygote-gradient-does-not-work-with-abstractgps-custommean/87815/2 "2022-09-26T14:34:05Z")

</div>

Hi @soldasim . This sounds like a frustrating issue.

Would I be correct in assuming that you’re not actually interested in using the custom mean function `CustomMean(x -> 0.)`? Would you mind providing the one that you’re actually interested in?

---

<div class="post-metadata">

**Author:** ![soldasim](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/soldasim/32/42514_2.png) [@soldasim](https://discourse.julialang.org/u/soldasim)\
**Post date:** [September 26, 2022, 2:45pm UTC](https://discourse.julialang.org/t/zygote-gradient-does-not-work-with-abstractgps-custommean/87815/3 "2022-09-26T14:45:26Z")

</div>

Hi @willtebbutt,

in my model I am using a parametric model as the mean of the GP.  
I want the code to be general, so the parametric model is to be defined by the user, so I need it to work with any real function.

An example could be:

```julia
function param_model(x, params)
    return params[1] * cos(params[2] * x[1]) + params[3]
end

zygote_custommean = Zygote.gradient(n -> f(n; mean=AbstractGPs.CustomMean(x -> param_model(x, [1.,2.,3.]))), 1.)

```

EDIT: error in the example

---

<div class="post-metadata">

**Author:** ![willtebbutt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/willtebbutt/32/6790_2.png) [@willtebbutt](https://discourse.julialang.org/u/willtebbutt)\
**Post date:** [September 26, 2022, 3:24pm UTC](https://discourse.julialang.org/t/zygote-gradient-does-not-work-with-abstractgps-custommean/87815/4 "2022-09-26T15:24:46Z")

</div>

Okay, thanks for the additional info.

One more question: is the actual value of `X` that you plan to use going to be multi-dimensional, or do you only ever do 1D stuff?

---

<div class="post-metadata">

**Author:** ![soldasim](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/soldasim/32/42514_2.png) [@soldasim](https://discourse.julialang.org/u/soldasim)\
**Post date:** [September 26, 2022, 3:26pm UTC](https://discourse.julialang.org/t/zygote-gradient-does-not-work-with-abstractgps-custommean/87815/5 "2022-09-26T15:26:12Z")

</div>

`X` can be multidimensional.

---

<div class="post-metadata">

**Author:** ![willtebbutt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/willtebbutt/32/6790_2.png) [@willtebbutt](https://discourse.julialang.org/u/willtebbutt)\
**Post date:** [September 26, 2022, 3:41pm UTC](https://discourse.julialang.org/t/zygote-gradient-does-not-work-with-abstractgps-custommean/87815/6 "2022-09-26T15:41:12Z")

</div>

Great.

In that case I would suggest implementing a new mean function, so that you can implement `_map_meanfunction` for it, and get the appropriate cotangent type when you work with `ColVecs` and `RowVecs`.

In particular, I would do something like

```julia
struct ParamModel{Tparams} <: AbstractGPs.MeanFunction
    params::Tparams
end

function AbstractGPs._map_meanfunction(m::ParamModel, x::ColVecs)
    X = x.X
    # whatever computation needs to happen on `X` using broadcasting etc
end

function AbstractGPs._map_meanfunction(m::ParamModel, x::RowVecs)
    X = x.X
    # whatever computation needs to happen on `X` using broadcasting etc
end

```

What’s gong wrong with AD at the minute is that the `ChainRulesCore.rrule` for `map` doesn’t do the right thing for `ColVecs` and `RowVecs`, which is what `_map_meanfunction` hits for `CustomMean`. By implementing `_map_meanfunction` on your own type, you can ensure that you never hit this code path.

It’s an annoying work-around to have to make, but that’s (unfortunately) where we are with AD at the minute.

---

<div class="post-metadata">

**Author:** ![soldasim](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/soldasim/32/42514_2.png) [@soldasim](https://discourse.julialang.org/u/soldasim)\
**Post date:** [September 26, 2022, 4:33pm UTC](https://discourse.julialang.org/t/zygote-gradient-does-not-work-with-abstractgps-custommean/87815/7 "2022-09-26T16:33:08Z")

</div>

Thank you so much! It would have been difficult for me to find the cause as I’m not that experienced in julia yet.

I’ve generalized the solution into

```julia
struct MyCustomMean{Tf} <: AbstractGPs.MeanFunction
    f::Tf
end

function AbstractGPs._map_meanfunction(m::MyCustomMean, x::ColVecs)
    X = x.X
    map(m.f, eachcol(X))
end

function AbstractGPs._map_meanfunction(m::MyCustomMean, x::RowVecs)
    X = x.X
    map(m.f, eachcol(X))
end

```

so that `MyCustomMean` can be used exactly like `AbstractGPs.CustomMean`.

The code seems to work now:

```julia
julia> ForwardDiff.gradient(n -> f(n; mean=AbstractGPs.CustomMean(x->0.)), [1.])
1-element Vector{Float64}:
 0.2630913490608584

julia> ForwardDiff.gradient(n -> f(n; mean=MyCustomMean(x->0.)), [1.])
1-element Vector{Float64}:
 0.2630913490608584

julia> Zygote.gradient(n -> f(n; mean=MyCustomMean(x->0.)), 1.)
(0.26309134906085807,)

```

Thank you for your time and help 🙂

---

<div class="post-metadata">

**Author:** ![kunal](https://avatars.discourse-cdn.com/v4/letter/k/dec6dc/32.png) [@kunal](https://discourse.julialang.org/u/kunal)\
**Post date:** [July 30, 2024, 6:54pm UTC](https://discourse.julialang.org/t/zygote-gradient-does-not-work-with-abstractgps-custommean/87815/8 "2024-07-30T18:54:57Z")

</div>

Hi, could you please let me know where I am doing this wrong? or something has changed.

```julia
using AbstractGPs
using Zygote
# using ForwardDiff

X = [1.; 2.; 3.;;]
y = [1., 2., 3.]

function construct_finite_gp(X, lengthscale, noise; min_param_val=1e-6, mean)
    # for numerical stability
    lengthscale = lengthscale + min_param_val
    noise = noise + min_param_val

    kernel = with_lengthscale(Matern52Kernel(), lengthscale)
    return GP(mean, kernel)(X', noise)
end

function f(noise; mean)
    gp = construct_finite_gp(X, 1., first(noise); mean)
    return logpdf(gp, y)
end

struct MyCustomMean{Tf} <: AbstractGPs.MeanFunction
    f::Tf
end

function AbstractGPs._map_meanfunction(m::MyCustomMean, x::ColVecs)
    X = x.X
    map(m.f, eachcol(X))
end

function AbstractGPs._map_meanfunction(m::MyCustomMean, x::RowVecs)
    X = x.X
    map(m.f, eachcol(X))
end

# zygote_zeromean = Zygote.gradient(n -> f(n; mean=AbstractGPs.ZeroMean()), 1.)
# @show zygote_zeromean

Zygote.gradient(n -> f(n; mean=MyCustomMean(x->0.)), 1.)

```

 ![image](https://global.discourse-cdn.com/julialang/original/3X/a/c/accf364a4f048b7def104fa4a12762409533701e.png)

---

<div class="post-metadata">

**Author:** ![soldasim](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/soldasim/32/42514_2.png) [@soldasim](https://discourse.julialang.org/u/soldasim)\
**Post date:** [July 31, 2024, 5:27pm UTC](https://discourse.julialang.org/t/zygote-gradient-does-not-work-with-abstractgps-custommean/87815/9 "2024-07-31T17:27:36Z")

</div>

I am not sure what changed but [this](https://github.com/soldasim/BOSS.jl/blob/master/src/models/gaussian_process.jl) is how I currently do it.

You don’t need any `MyCustomMean` structure, just pass the custom mean function as a function. The ForwardDiff works as is, and you need these two methods as a workaround for zygote to work;

```julia
AbstractGPs.mean_vector(m::AbstractGPs.CustomMean, x::ColVecs) = map(m.f, eachcol(x.X))
AbstractGPs.mean_vector(m::AbstractGPs.CustomMean, x::RowVecs) = map(m.f, eachrow(x.X))

```

The complete fixed script can look like this

```julia
using AbstractGPs
using Zygote
using ForwardDiff

X = [1.;; 2.;; 3.;;]
y = [1., 2., 3.]

function construct_finite_gp(X, lengthscale, noise; min_param_val=1e-6, mean)
    # for numerical stability
    lengthscale = lengthscale + min_param_val
    noise = noise + min_param_val

    kernel = with_lengthscale(Matern52Kernel(), lengthscale)
    return GP(mean, kernel)(X, noise)
end

function f(noise; mean)
    gp = construct_finite_gp(X, 1., first(noise); mean)
    return logpdf(gp, y)
end

# workaround needed for Zygote
AbstractGPs.mean_vector(m::AbstractGPs.CustomMean, x::ColVecs) = map(m.f, eachcol(x.X))
AbstractGPs.mean_vector(m::AbstractGPs.CustomMean, x::RowVecs) = map(m.f, eachrow(x.X))

zygote = Zygote.gradient(n -> f(n; mean=x->0.), 1.)
fwddiff = ForwardDiff.gradient(n -> f(n; mean=x->0.), [1.])

@show zygote
@show fwddiff

```

---

<div class="post-metadata">

**Author:** ![kunal](https://avatars.discourse-cdn.com/v4/letter/k/dec6dc/32.png) [@kunal](https://discourse.julialang.org/u/kunal)\
**Post date:** [August 1, 2024, 6:06am UTC](https://discourse.julialang.org/t/zygote-gradient-does-not-work-with-abstractgps-custommean/87815/10 "2024-08-01T06:06:37Z")

</div>

Thanks a lot @soldasim!  
This worked for me.  
Unfortunately, AbstractGPs doesn’t adapt with Zygote quickly. I think they need to vectorize the mean operation.