Zygote.gradient does not work with AbstractGPs.CustomMean

General Usage
, ,
3

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:

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