fix MaternKernel AD, but remove differentiation wrt \nu (#425)

* reactivate test_ADs of MaternKernel, but do not test differentiation through \nu argument
* make clear in docstring that MaternKernel does not support derivative w.r.t. order \nu
* work-around for Zygote AD that will return nothing/NaN gradients w.r.t. \nu but fixes the gradient w.r.t. x

Author: st--

Project.toml

@@ -1,6 +1,6 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.10.36"
+version = "0.10.37"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/basekernels/matern.jl

@@ -17,6 +17,10 @@ By default, ``d`` is the Euclidean metric ``d(x, x') = \\|x - x'\\|_2``.
 A Gaussian process with a Matérn kernel is ``\\lceil \\nu \\rceil - 1``-times
 differentiable in the mean-square sense.
 
+!!! note
+
+    Differentiation with respect to the order ν is not currently supported.
+
 See also: [`Matern12Kernel`](@ref), [`Matern32Kernel`](@ref), [`Matern52Kernel`](@ref)
 """
 struct MaternKernel{Tν<:Real,M} <: SimpleKernel
@@ -33,8 +37,11 @@ MaternKernel(; nu::Real=1.5, ν::Real=nu, metric=Euclidean()) = MaternKernel(ν,
 
 @functor MaternKernel
 
-@inline function kappa(κ::MaternKernel, d::Real)
-    result = _matern(only(κ.ν), d)
+@inline _get_ν(k::MaternKernel) = only(k.ν)
+ChainRulesCore.@non_differentiable _get_ν(k)  # work-around; should be "NotImplemented" rather than NoTangent
+
+@inline function kappa(k::MaternKernel, d::Real)
+    result = _matern(_get_ν(k), d)
     return ifelse(iszero(d), one(result), result)
 end
 

src/matrix/kernelkroneckermat.jl

@@ -14,7 +14,7 @@ where `D` is given by `dims`.
 
 !!! warning
 
-    Require `Kronecker.jl` and for `iskroncompatible(κ)` to return `true`.
+    Requires `Kronecker.jl` and for `iskroncompatible(κ)` to return `true`.
 """
 function kernelkronmat(κ::Kernel, X::AbstractVector{<:Real}, dims::Int)
     checkkroncompatible(κ)

test/basekernels/matern.jl

@@ -4,25 +4,24 @@
     v1 = rand(rng, 3)
     v2 = rand(rng, 3)
     @testset "MaternKernel" begin
-        ν = 2.0
+        ν = 2.1
         k = MaternKernel(; ν=ν)
         matern(x, ν) = 2^(1 - ν) / gamma(ν) * (sqrt(2ν) * x)^ν * besselk(ν, sqrt(2ν) * x)
         @test MaternKernel(; nu=ν).ν == [ν]
         @test kappa(k, x) ≈ matern(x, ν)
         @test kappa(k, 0.0) == 1.0
-        @test kappa(MaternKernel(; ν=ν), x) == kappa(k, x)
         @test metric(MaternKernel()) == Euclidean()
         @test metric(MaternKernel(; ν=2.0)) == Euclidean()
         @test repr(k) == "Matern Kernel (ν = $(ν), metric = Euclidean(0.0))"
-        # test_ADs(x->MaternKernel(nu=first(x)),[ν])
-        @test_broken "All fails (because of logabsgamma for ForwardDiff and ReverseDiff and because of nu for Zygote)"
 
         k2 = MaternKernel(; ν=ν, metric=WeightedEuclidean(ones(3)))
         @test metric(k2) isa WeightedEuclidean
         @test k2(v1, v2) ≈ k(v1, v2)
 
         # Standardised tests.
         TestUtils.test_interface(k, Float64)
+        test_ADs(() -> MaternKernel(; nu=ν))
+
         test_params(k, ([ν],))
     end
     @testset "Matern32Kernel" begin