Example code in README fails (on Zygote.gradient)

[theo-brown]

Hi all,

The autodiff part of the example code in the README fails. If the Zygote line is removed, then it runs fine, but with it, you get the following stack trace:

julia> training_results = Optim.optimize(
           objective ∘ unflatten,
           θ -> only(Zygote.gradient(objective ∘ unflatten, θ)),
           flat_initial_params,
           BFGS(
               alphaguess = Optim.LineSearches.InitialStatic(scaled=true),
               linesearch = Optim.LineSearches.BackTracking(),
           ),
           Optim.Options(show_trace = true);
           inplace=false,
       )

ERROR: MethodError: no method matching AbstractGPs.FiniteGP(::GP{AbstractGPs.ZeroMean{Float64}, KernelSum{Tuple{TransformedKernel{ScaledKernel{Matern52Kernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}, TransformedKernel{ScaledKernel{SqExponentialKernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}}}}, ::Float64)

Closest candidates are:
  AbstractGPs.FiniteGP(::AbstractGPs.AbstractGP, ::AbstractVector)
   @ AbstractGPs ~/.julia/packages/AbstractGPs/XejGR/src/finite_gp_projection.jl:19
  AbstractGPs.FiniteGP(::AbstractGPs.AbstractGP, ::AbstractVector, ::Real)
   @ AbstractGPs ~/.julia/packages/AbstractGPs/XejGR/src/finite_gp_projection.jl:19
  AbstractGPs.FiniteGP(::AbstractGPs.AbstractGP, ::AbstractVector, ::AbstractVector{<:Real})
   @ AbstractGPs ~/.julia/packages/AbstractGPs/XejGR/src/finite_gp_projection.jl:13
  ...

Stacktrace:
  [1] macro expansion
    @ ~/.julia/packages/Zygote/YYT6v/src/compiler/interface2.jl:101 [inlined]
  [2] _pullback(::Zygote.Context{false}, ::Type{AbstractGPs.FiniteGP}, ::GP{AbstractGPs.ZeroMean{Float64}, KernelSum{Tuple{TransformedKernel{ScaledKernel{Matern52Kernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}, TransformedKernel{ScaledKernel{SqExponentialKernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}}}}, ::Float64)
    @ Zygote ~/.julia/packages/Zygote/YYT6v/src/compiler/interface2.jl:101
  [3] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
  [4] adjoint
    @ ~/.julia/packages/Zygote/YYT6v/src/lib/lib.jl:203 [inlined]
  [5] _pullback
    @ ~/.julia/packages/ZygoteRules/OgCVT/src/adjoint.jl:66 [inlined]
  [6] _pullback
    @ ~/.julia/packages/AbstractGPs/XejGR/src/finite_gp_projection.jl:32 [inlined]
  [7] _pullback(ctx::Zygote.Context{false}, f::GP{AbstractGPs.ZeroMean{Float64}, KernelSum{Tuple{TransformedKernel{ScaledKernel{Matern52Kernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}, TransformedKernel{ScaledKernel{SqExponentialKernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}}}}, args::Float64)
    @ Zygote ~/.julia/packages/Zygote/YYT6v/src/compiler/interface2.jl:0
  [8] #rrule_via_ad#54
    @ ~/.julia/packages/Zygote/YYT6v/src/compiler/chainrules.jl:260 [inlined]
  [9] rrule_via_ad
    @ ~/.julia/packages/Zygote/YYT6v/src/compiler/chainrules.jl:248 [inlined]
 [10] #1659
    @ ./none:0 [inlined]
 [11] iterate
    @ ./generator.jl:47 [inlined]
 [12] collect(itr::Base.Generator{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, ChainRules.var"#1659#1664"{Zygote.ZygoteRuleConfig{Zygote.Context{false}}, GP{AbstractGPs.ZeroMean{Float64}, KernelSum{Tuple{TransformedKernel{ScaledKernel{Matern52Kernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}, TransformedKernel{ScaledKernel{SqExponentialKernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}}}}}})
    @ Base ./array.jl:782
 [13] rrule(config::Zygote.ZygoteRuleConfig{Zygote.Context{false}}, ::typeof(sum), f::GP{AbstractGPs.ZeroMean{Float64}, KernelSum{Tuple{TransformedKernel{ScaledKernel{Matern52Kernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}, TransformedKernel{ScaledKernel{SqExponentialKernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}}}}, xs::StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}; dims::Function)
    @ ChainRules ~/.julia/packages/ChainRules/snrkz/src/rulesets/Base/mapreduce.jl:102
 [14] rrule
    @ ~/.julia/packages/ChainRules/snrkz/src/rulesets/Base/mapreduce.jl:76 [inlined]
 [15] rrule(config::Zygote.ZygoteRuleConfig{Zygote.Context{false}}, ::typeof(mean), f::GP{AbstractGPs.ZeroMean{Float64}, KernelSum{Tuple{TransformedKernel{ScaledKernel{Matern52Kernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}, TransformedKernel{ScaledKernel{SqExponentialKernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}}}}, x::StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}; dims::Function)
    @ ChainRules ~/.julia/packages/ChainRules/snrkz/src/rulesets/Statistics/statistics.jl:28
 [16] rrule
    @ ~/.julia/packages/ChainRules/snrkz/src/rulesets/Statistics/statistics.jl:21 [inlined]
 [17] chain_rrule
    @ ~/.julia/packages/Zygote/YYT6v/src/compiler/chainrules.jl:223 [inlined]
 [18] macro expansion
    @ ~/.julia/packages/Zygote/YYT6v/src/compiler/interface2.jl:101 [inlined]
 [19] _pullback(::Zygote.Context{false}, ::typeof(mean), ::GP{AbstractGPs.ZeroMean{Float64}, KernelSum{Tuple{TransformedKernel{ScaledKernel{Matern52Kernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}, TransformedKernel{ScaledKernel{SqExponentialKernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}}}}, ::StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64})
    @ Zygote ~/.julia/packages/Zygote/YYT6v/src/compiler/interface2.jl:101
 [20] _pullback
    @ ~/.julia/packages/AbstractGPs/XejGR/src/abstract_gp.jl:48 [inlined]
 [21] _pullback(::Zygote.Context{false}, ::typeof(mean_and_cov), ::GP{AbstractGPs.ZeroMean{Float64}, KernelSum{Tuple{TransformedKernel{ScaledKernel{Matern52Kernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}, TransformedKernel{ScaledKernel{SqExponentialKernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}}}}, ::StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64})
    @ Zygote ~/.julia/packages/Zygote/YYT6v/src/compiler/interface2.jl:0
 [22] _pullback
    @ ~/.julia/packages/AbstractGPs/XejGR/src/finite_gp_projection.jl:134 [inlined]
 [23] _pullback(ctx::Zygote.Context{false}, f::typeof(mean_and_cov), args::AbstractGPs.FiniteGP{GP{AbstractGPs.ZeroMean{Float64}, KernelSum{Tuple{TransformedKernel{ScaledKernel{Matern52Kernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}, TransformedKernel{ScaledKernel{SqExponentialKernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}}}}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, LinearAlgebra.Diagonal{Float64, FillArrays.Fill{Float64, 1, Tuple{Base.OneTo{Int64}}}}})
    @ Zygote ~/.julia/packages/Zygote/YYT6v/src/compiler/interface2.jl:0
 [24] _pullback
    @ ~/.julia/packages/AbstractGPs/XejGR/src/finite_gp_projection.jl:307 [inlined]
 [25] _pullback(::Zygote.Context{false}, ::typeof(logpdf), ::AbstractGPs.FiniteGP{GP{AbstractGPs.ZeroMean{Float64}, KernelSum{Tuple{TransformedKernel{ScaledKernel{Matern52Kernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}, TransformedKernel{ScaledKernel{SqExponentialKernel{Distances.Euclidean}, Float64}, ScaleTransform{Float64}}}}}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, LinearAlgebra.Diagonal{Float64, FillArrays.Fill{Float64, 1, Tuple{Base.OneTo{Int64}}}}}, ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/YYT6v/src/compiler/interface2.jl:0
 [26] _pullback
    @ ./REPL[29]:4 [inlined]
 [27] _pullback(ctx::Zygote.Context{false}, f::typeof(objective), args::NamedTuple{(:k1, :k2, :noise_var), Tuple{NamedTuple{(:var, :precision), Tuple{Float64, Float64}}, NamedTuple{(:var, :precision), Tuple{Float64, Float64}}, Float64}})
    @ Zygote ~/.julia/packages/Zygote/YYT6v/src/compiler/interface2.jl:0
 [28] _pullback
    @ ./operators.jl:1034 [inlined]
 [29] _pullback
    @ ./operators.jl:1031 [inlined]
 [30] _pullback(::Zygote.Context{false}, ::Base.var"##_#97", ::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::ComposedFunction{typeof(objective), ComposedFunction{typeof(ParameterHandling.value), ParameterHandling.var"#unflatten_to_NamedTuple#17"{Float64, NamedTuple{(:k1, :k2, :noise_var), Tuple{NamedTuple{(:var, :precision), Tuple{ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, NamedTuple{(:var, :precision), Tuple{ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, ParameterHandling.var"#unflatten_to_Tuple#15"{Float64, Tuple{Int64, Int64, Int64}, Tuple{Int64, Int64, Int64}, Tuple{ParameterHandling.var"#unflatten_to_NamedTuple#17"{Float64, NamedTuple{(:var, :precision), Tuple{ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, ParameterHandling.var"#unflatten_to_Tuple#15"{Float64, Tuple{Int64, Int64}, Tuple{Int64, Int64}, Tuple{ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}, ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}}}}, ParameterHandling.var"#unflatten_to_NamedTuple#17"{Float64, NamedTuple{(:var, :precision), Tuple{ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, ParameterHandling.var"#unflatten_to_Tuple#15"{Float64, Tuple{Int64, Int64}, Tuple{Int64, Int64}, Tuple{ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}, ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}}}}, ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}}}}}}, ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/YYT6v/src/compiler/interface2.jl:0
 [31] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [32] adjoint
    @ ~/.julia/packages/Zygote/YYT6v/src/lib/lib.jl:203 [inlined]
 [33] _pullback
    @ ~/.julia/packages/ZygoteRules/OgCVT/src/adjoint.jl:66 [inlined]
 [34] _pullback
    @ ./operators.jl:1031 [inlined]
 [35] _pullback(ctx::Zygote.Context{false}, f::ComposedFunction{typeof(objective), ComposedFunction{typeof(ParameterHandling.value), ParameterHandling.var"#unflatten_to_NamedTuple#17"{Float64, NamedTuple{(:k1, :k2, :noise_var), Tuple{NamedTuple{(:var, :precision), Tuple{ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, NamedTuple{(:var, :precision), Tuple{ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, ParameterHandling.var"#unflatten_to_Tuple#15"{Float64, Tuple{Int64, Int64, Int64}, Tuple{Int64, Int64, Int64}, Tuple{ParameterHandling.var"#unflatten_to_NamedTuple#17"{Float64, NamedTuple{(:var, :precision), Tuple{ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, ParameterHandling.var"#unflatten_to_Tuple#15"{Float64, Tuple{Int64, Int64}, Tuple{Int64, Int64}, Tuple{ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}, ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}}}}, ParameterHandling.var"#unflatten_to_NamedTuple#17"{Float64, NamedTuple{(:var, :precision), Tuple{ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, ParameterHandling.var"#unflatten_to_Tuple#15"{Float64, Tuple{Int64, Int64}, Tuple{Int64, Int64}, Tuple{ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}, ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}}}}, ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}}}}}}, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/YYT6v/src/compiler/interface2.jl:0
 [36] pullback(f::Function, cx::Zygote.Context{false}, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/YYT6v/src/compiler/interface.jl:44
 [37] pullback
    @ ~/.julia/packages/Zygote/YYT6v/src/compiler/interface.jl:42 [inlined]
 [38] gradient(f::Function, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/YYT6v/src/compiler/interface.jl:96
 [39] (::var"#3#4")(θ::Vector{Float64})
    @ Main ./REPL[30]:9
 [40] (::NLSolversBase.var"#gg!#2"{var"#3#4"})(G::Vector{Float64}, x::Vector{Float64})
    @ NLSolversBase ~/.julia/packages/NLSolversBase/kavn7/src/objective_types/inplace_factory.jl:21
 [41] (::NLSolversBase.var"#fg!#8"{ComposedFunction{typeof(objective), ComposedFunction{typeof(ParameterHandling.value), ParameterHandling.var"#unflatten_to_NamedTuple#17"{Float64, NamedTuple{(:k1, :k2, :noise_var), Tuple{NamedTuple{(:var, :precision), Tuple{ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, NamedTuple{(:var, :precision), Tuple{ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, ParameterHandling.var"#unflatten_to_Tuple#15"{Float64, Tuple{Int64, Int64, Int64}, Tuple{Int64, Int64, Int64}, Tuple{ParameterHandling.var"#unflatten_to_NamedTuple#17"{Float64, NamedTuple{(:var, :precision), Tuple{ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, ParameterHandling.var"#unflatten_to_Tuple#15"{Float64, Tuple{Int64, Int64}, Tuple{Int64, Int64}, Tuple{ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}, ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}}}}, ParameterHandling.var"#unflatten_to_NamedTuple#17"{Float64, NamedTuple{(:var, :precision), Tuple{ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.Positive{Float64, typeof(exp), Float64}}}, ParameterHandling.var"#unflatten_to_Tuple#15"{Float64, Tuple{Int64, Int64}, Tuple{Int64, Int64}, Tuple{ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}, ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}}}}, ParameterHandling.var"#unflatten_Positive#25"{Float64, ParameterHandling.Positive{Float64, typeof(exp), Float64}, ParameterHandling.var"#unflatten_to_Real#2"{Float64, Float64}}}}}}}, NLSolversBase.var"#gg!#2"{var"#3#4"}})(gx::Vector{Float64}, x::Vector{Float64})
    @ NLSolversBase ~/.julia/packages/NLSolversBase/kavn7/src/objective_types/abstract.jl:13
 [42] value_gradient!!(obj::OnceDifferentiable{Float64, Vector{Float64}, Vector{Float64}}, x::Vector{Float64})
    @ NLSolversBase ~/.julia/packages/NLSolversBase/kavn7/src/interface.jl:82
 [43] initial_state(method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.BackTracking{Float64, Int64}, Nothing, Nothing, Flat}, options::Optim.Options{Float64, Nothing}, d::OnceDifferentiable{Float64, Vector{Float64}, Vector{Float64}}, initial_x::Vector{Float64})
    @ Optim ~/.julia/packages/Optim/V8ZEC/src/multivariate/solvers/first_order/bfgs.jl:94
 [44] optimize
    @ ~/.julia/packages/Optim/V8ZEC/src/multivariate/optimize/optimize.jl:36 [inlined]
 [45] optimize(f::Function, g::Function, initial_x::Vector{Float64}, method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.BackTracking{Float64, Int64}, Nothing, Nothing, Flat}, options::Optim.Options{Float64, Nothing}; inplace::Bool, autodiff::Symbol)
    @ Optim ~/.julia/packages/Optim/V8ZEC/src/multivariate/optimize/interface.jl:156
 [46] top-level scope
    @ REPL[30]:7

Version info:

julia> versioninfo()
Julia Version 1.9.3
Commit bed2cd540a1 (2023-08-24 14:43 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 8 × 11th Gen Intel(R) Core(TM) i5-1135G7 @ 2.40GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, tigerlake)
  Threads: 1 on 8 virtual cores

[willtebbutt]

Interesting. I seems to work for me on a clean install. Do you have other packages installed / could you see you see whether you can get it to work on a clean install?

[theo-brown]

Oh god not me screwing up my environment again

My packages:

> Pkg.status()
Status `~/.julia/environments/v1.9/Project.toml`
⌃ [99985d1d] AbstractGPs v0.5.9
  [861a8166] Combinatorics v1.0.2
  [b4f34e82] Distances v0.10.10
  [31c24e10] Distributions v0.25.102
  [f6369f11] ForwardDiff v0.10.36
  [891a1506] GaussianProcesses v0.12.5
  [ec8451be] KernelFunctions v0.10.58
  [429524aa] Optim v1.7.8
  [2412ca09] ParameterHandling v0.4.7
  [91a5bcdd] Plots v1.39.0
  [e88e6eb3] Zygote v0.6.67
  [37e2e46d] LinearAlgebra

[willtebbutt]

Ah. It looks ike you've got a really old version of AbstractGPs installed -- 0.5.9 rather than the most recent patch version,0.5.19. Does Pkg.update() give you the newest version?

[theo-brown]

I think it was pinned to 0.5.9 by GaussianProcesses.jl, but I'll need to check that.

[willtebbutt]

Ah okay. We don't check for compatibility with GaussianProcesses.jl because our set of packages are completely separate. If you want to use both of them, you might need to work in separate environments.

[theo-brown]

Yep, I've fully switched over to AbstractGPs anyway now, I just hadn't removed it. Uninstalling GaussianProcesses.jl and upgrading AbstractGPs fixed it, thanks.