Known AD Failures

[theogf]

Here is a list of the failures in the tests made in #114 I observed with the different ADs : ForwardDiff.jl, Zygote.jl and ReverseDiff.jl :

• FBMKernel : ForwardDiff/Zygote
• GaborKernel : Zygote (ForwardDiff+ReverseDiff randomly)
• MaternKernel : Zygote
• SpectralMixture : All
• PeriodicKernel : Zygote (need to define an adjoint) Fix periodic kernel AD #531, ForwardDiff and ReverseDiff fail randomly (edge cases might need to be looked at) PeriodicKernel does not work with AD (see issue for work-around) #389
• Wiener : All
• GammaExponentialKernel : Zygote
• NeuralNetworkKernel : Zygote
• PiecewisePolynomialKernel : All
• PolynomialKernel : All
• GammaRationalQuadratic : All
• KernelProduct : Zygote
• KernelSum : Zygote
• FunctionTransform : Zygote
• ChainTransform : Zygote Zygote AD failure for ChainTransform #415

This is a good starting point to try to find solutions

[devmotion]

The Zygote problems with MaternKernel are caused by the fact that the partial derivative of besselk with respect to the first argument is defined as NaN in https://github.com/JuliaDiff/ChainRules.jl/blob/98c54587257b86cce6eb45f7870a75f897058d21/src/rulesets/packages/SpecialFunctions.jl#L46-L47 (and I assume the same problem exists for the other AD backends, since I get NaN for all of them when I try to run the commented out AD tests). I guess one would have to implement https://dlmf.nist.gov/10.38 to fix it.

[theogf]

Haha writing these derivatives sounds like one should write a whole package about bessel functions

[yebai]

@sharanry Can you prioritise these AD issues? It would be great if these issues can be addressed during the summer.

[devmotion]

BTW I found some publication from 2016 with closed-form expressions of the derivatives of the Bessel functions with respect to the order. I opened an issue at JuliaDiff/ChainRules.jl#208 to discuss how one would deal with the additional dependencies needed for their implementations (they contain hypergeometric functions).

[devmotion]

We might want to refactor KernelSum and KernelProd (making them concretely typed and allowing both tuples and vectors of kernels similar to TensorProduct, and probably removing the weights in KernelSum) before fixing any AD issues there.

[theogf]

Agreed!
There is also a general AD issue when using Transform where the pullback on ColVecs and RowVecs return a vector of vectors, this would tick off a good portions of the issues.

[sharanry]

@sharanry Can you prioritise these AD issues? It would be great if these issues can be addressed during the summer.

Sorry for the late reply. I somehow didn't get a notification this comment. Randomly found this while browsing the issues. I am looking into it.

[sharanry]

The probable reason Zygote fails for FunctionTransform is the usage of Base.mapslices in

KernelFunctions.jl/src/transform/functiontransform.jl

Lines 19 to 20 in 0df9e83

	_map(t::FunctionTransform, x::ColVecs) = ColVecs(mapslices(t.f, x.X; dims=1))
	_map(t::FunctionTransform, x::RowVecs) = RowVecs(mapslices(t.f, x.X; dims=2))

.

Base.mapslices is mutating the array. Not sure why. JuliaLang/julia#17266 should have fixed this.

julia> Zygote._pullback(x-> mapslices(x->sin.(x), x, dims=1), rand(3,3))[2](ones(3,3))
ERROR: Mutating arrays is not supported

[devmotion]

mapslices still mutates a temporary array. The linked PR just ensures that

mapslices never modifies the input array. It allocates temporary storage and copies each slice into it before calling the user-function.

[sharanry]

mapslices still mutates a temporary array. The linked PR just ensures that

Oh makes sense. Do you see any other efficient way to apply a function transform for a matrix/ColVecs/RowVecs?