Parametrized distances and Zygote

I am using Distances.jl for kernel functions, and have started looking into Zygote for some automatic differentiation. However, for a parametrized distance like WeightedSqEuclidean this runs into a problem.

julia> using Distances, Zygote
julia> x1 = rand(2); x2 = rand(2);
julia> gradient(x1 -> evaluate(WeightedSqEuclidean([1.0, 1.0]), x1, x2), x1)

Can't differentiate loopinfo expression

Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] macro expansion at ./simdloop.jl:79 [inlined]
 [3] _evaluate at /home/steffen/.julia/packages/Distances/jwhuc/src/wmetrics.jl:65 [inlined]
 [4] WeightedSqEuclidean at /home/steffen/.julia/packages/Distances/jwhuc/src/wmetrics.jl:84 [inlined]
 [5] (::typeof(∂(λ)))(::Float64) at /home/steffen/.julia/packages/Zygote/seGHk/src/compiler/interface2.jl:0
 [6] evaluate at /home/steffen/.julia/packages/Distances/jwhuc/src/generic.jl:24 [inlined]
 [7] #12 at ./In[20]:1 [inlined]
 [8] (::typeof(∂(#12)))(::Float64) at /home/steffen/.julia/packages/Zygote/seGHk/src/compiler/interface2.jl:0
 [9] (::Zygote.var"#41#42"{typeof(∂(#12))})(::Float64) at /home/steffen/.julia/packages/Zygote/seGHk/src/compiler/interface.jl:45
 [10] gradient(::Function, ::Array{Float64,1}) at /home/steffen/.julia/packages/Zygote/seGHk/src/compiler/interface.jl:54
 [11] top-level scope at In[20]:1

At the same time, the non-parametrized euclidean distances works well

julia> gradient(x1 -> evaluate(SqEuclidean(), x1, x2), x1)

([1.7071687253858214, 0.9053498617738835],)

I know that KernelFunctions.jl uses transforms on non-parametrized kernels, but I was wondering if it might be worthwhile to make parametrized distances play nice with Zygote. That being said, while I am not really clear what the issue is, what I have been skim-reading implies that making such a change might come with performance regressions. If that is the case it might not be worth it.

JuliaStats / Distances.jl

Parametrized distances and Zygote #171