Remove workaround

Author: devmotion

src/basekernels/matern.jl

@@ -37,21 +37,8 @@ MaternKernel(; nu::Real=1.5, ν::Real=nu, metric=Euclidean()) = MaternKernel(ν,
 
 @functor MaternKernel
 
-# workaround for Zygote
-# unclear why it's needed but it is fine since it's stated officially that we don't support differentiation with respect to ν
-@inline _get_ν(k::MaternKernel) = only(k.ν)
-function ChainRulesCore.rrule(::typeof(_get_ν), k::T) where {T<:MaternKernel}
-    function _get_ν_pullback(Δ)
-        dν = ChainRulesCore.@not_implemented(
-            "derivatives of `MaternKernel` w.r.t. order `ν` are not implemented."
-        )
-        return NoTangent(), Tangent{T}(; ν=dν, metric=NoTangent())
-    end
-    return _get_ν(k), _get_ν_pullback
-end
-
 @inline function kappa(k::MaternKernel, d::Real)
-    result = _matern(_get_ν(k), d)
+    result = _matern(only(k.ν), d)
     return ifelse(iszero(d), one(result), result)
 end
 

test/Project.toml

@@ -1,6 +1,5 @@
 [deps]
 AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
-ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
@@ -20,7 +19,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 AxisArrays = "0.4.3"
-ChainRulesTestUtils = "1.7"
 Compat = "3"
 Distances = "0.10"
 Documenter = "0.25, 0.26, 0.27"
@@ -32,4 +30,4 @@ LogExpFunctions = "0.2, 0.3"
 PDMats = "0.9, 0.10, 0.11"
 ReverseDiff = "1.2"
 SpecialFunctions = "0.10, 1, 2"
-Zygote = "0.4, 0.5, 0.6"
+Zygote = "0.6.38"

test/basekernels/matern.jl

@@ -18,15 +18,6 @@
         @test metric(k2) isa WeightedEuclidean
         @test k2(v1, v2) ≈ k(v1, v2)
 
-        # Test custom `rrule` (Zygote workaround).
-        k = MaternKernel(; ν=rand())
-        test_rrule(
-            KernelFunctions._get_ν,
-            k ⊢ ChainRulesTestUtils.Tangent{typeof(k)}(;
-                ν=randn(), metric=ChainRulesTestUtils.NoTangent()
-            ),
-        )
-
         # Standardised tests.
         TestUtils.test_interface(k, Float64)
         test_ADs(() -> MaternKernel(; nu=ν))

test/runtests.jl

@@ -1,6 +1,5 @@
 using KernelFunctions
 using AxisArrays
-using ChainRulesTestUtils
 using Distances
 using Documenter
 using Functors: functor