Merge pull request #157 from devmotion/functors

Author: devmotion

Project.toml

@@ -1,10 +1,11 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.6.1"
+version = "0.7.0"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
@@ -16,6 +17,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 [compat]
 Compat = "2.2, 3"
 Distances = "0.9"
+Functors = "0.1"
 Requires = "1.0.1"
 SpecialFunctions = "0.8, 0.9, 0.10"
 StatsBase = "0.32, 0.33"

docs/src/create_kernel.md

@@ -33,10 +33,43 @@ Note that `BaseKernel` do not use `Distances.jl` and can therefore be a bit slow
 ### Additional Options
 
 Finally there are additional functions you can define to bring in more features:
- - `KernelFunctions.trainable(k::MyKernel)`: it defines the trainable parameters of your kernel, it should return a `Tuple` of your parameters.
-These parameters will be passed to the `Flux.params` function. For some examples see the `trainable.jl` file in `src/`
  - `KernelFunctions.iskroncompatible(k::MyKernel)`: if your kernel factorizes in dimensions, you can declare your kernel as `iskroncompatible(k) = true` to use Kronecker methods.
  - `KernelFunctions.dim(x::MyDataType)`: by default the dimension of the inputs will only be checked for vectors of type `AbstractVector{<:Real}`. If you want to check the dimensionality of your inputs, dispatch the `dim` function on your datatype. Note that `0` is the default.
  - `dim` is called within `KernelFunctions.validate_inputs(x::MyDataType, y::MyDataType)`, which can instead be directly overloaded if you want to run special checks for your input types.
  - `kernelmatrix(k::MyKernel, ...)`: you can redefine the diverse `kernelmatrix` functions to eventually optimize the computations.
  - `Base.print(io::IO, k::MyKernel)`: if you want to specialize the printing of your kernel
+
+KernelFunctions uses [Functors.jl](https://github.com/FluxML/Functors.jl) for specifying trainable kernel parameters
+in a way that is compatible with the [Flux ML framework](https://github.com/FluxML/Flux.jl).
+You can use `Functors.@functor` if all fields of your kernel struct are trainable. Note that optimization algorithms
+in Flux are not compatible with scalar parameters (yet), and hence vector-valued parameters should be preferred.
+
+```julia
+import Functors
+
+struct MyKernel{T} <: KernelFunctions.Kernel
+    a::Vector{T}
+end
+
+Functors.@functor MyKernel
+```
+
+If only a subset of the fields are trainable, you have to specify explicitly how to (re)construct the kernel with
+modified parameter values by [implementing `Functors.functor(::Type{<:MyKernel}, x)` for your kernel struct](https://github.com/FluxML/Functors.jl/issues/3):
+
+```julia
+import Functors
+
+struct MyKernel{T} <: KernelFunctions.Kernel
+    n::Int
+    a::Vector{T}
+end
+
+function Functors.functor(::Type{<:MyKernel}, x::MyKernel)
+    function reconstruct_mykernel(xs)
+        # keep field `n` of the original kernel and set `a` to (possibly different) `xs.a`
+        return MyKernel(x.n, xs.a)
+    end
+    return (a = x.a,), reconstruct_mykernel
+end
+```

src/KernelFunctions.jl

@@ -37,6 +37,7 @@ export IndependentMOKernel
 using Compat
 using Requires
 using Distances, LinearAlgebra
+using Functors
 using SpecialFunctions: loggamma, besselk, polygamma
 using ZygoteRules: @adjoint, pullback
 using StatsFuns: logtwo
@@ -79,7 +80,6 @@ include("zygote_adjoints.jl")
 function __init__()
     @require Kronecker="2c470bb0-bcc8-11e8-3dad-c9649493f05e" include("matrix/kernelkroneckermat.jl")
     @require PDMats="90014a1f-27ba-587c-ab20-58faa44d9150" include("matrix/kernelpdmat.jl")
-    @require Flux="587475ba-b771-5e3f-ad9e-33799f191a9c" include("trainable.jl")
 end
 
 end

src/basekernels/constant.jl

@@ -15,7 +15,6 @@ metric(::ZeroKernel) = Delta()
 
 Base.show(io::IO, ::ZeroKernel) = print(io, "Zero Kernel")
 
-
 """
     WhiteKernel()
 
@@ -55,6 +54,8 @@ struct ConstantKernel{Tc<:Real} <: SimpleKernel
     end
 end
 
+@functor ConstantKernel
+
 kappa(κ::ConstantKernel,x::Real) = first(κ.c)*one(x)
 
 metric(::ConstantKernel) = Delta()

src/basekernels/exponential.jl

@@ -63,6 +63,8 @@ struct GammaExponentialKernel{Tγ<:Real} <: SimpleKernel
     end
 end
 
+@functor GammaExponentialKernel
+
 kappa(κ::GammaExponentialKernel, d²::Real) = exp(-d²^first(κ.γ))
 
 metric(::GammaExponentialKernel) = SqEuclidean()

src/basekernels/fbm.jl

@@ -17,6 +17,8 @@ struct FBMKernel{T<:Real} <: Kernel
     end
 end
 
+@functor FBMKernel
+
 function (κ::FBMKernel)(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})
     modX = sum(abs2, x)
     modY = sum(abs2, y)

src/basekernels/gabor.jl

@@ -15,6 +15,8 @@ struct GaborKernel{K<:Kernel} <: Kernel
     end
 end
 
+@functor GaborKernel
+
 (κ::GaborKernel)(x, y) = κ.kernel(x ,y)
 
 function _gabor(; ell = nothing, p = nothing)

src/basekernels/maha.jl

@@ -16,6 +16,8 @@ struct MahalanobisKernel{T<:Real, A<:AbstractMatrix{T}} <: SimpleKernel
     end
 end
 
+@functor MahalanobisKernel
+
 kappa(κ::MahalanobisKernel, d::T) where {T<:Real} = exp(-d)
 
 metric(κ::MahalanobisKernel) = SqMahalanobis(κ.P)

src/basekernels/matern.jl

@@ -15,6 +15,8 @@ struct MaternKernel{Tν<:Real} <: SimpleKernel
     end
 end
 
+@functor MaternKernel
+
 @inline function kappa(κ::MaternKernel, d::Real)
     result = _matern(first(κ.ν), d)
     return ifelse(iszero(d), one(result), result)

src/basekernels/periodic.jl

@@ -20,6 +20,8 @@ PeriodicKernel(dims::Int) = PeriodicKernel(Float64, dims)
 
 PeriodicKernel(T::DataType, dims::Int = 1) = PeriodicKernel(r = ones(T, dims))
 
+@functor PeriodicKernel
+
 metric(κ::PeriodicKernel) = Sinus(κ.r)
 
 kappa(κ::PeriodicKernel, d::Real) = exp(- 0.5d)

src/basekernels/piecewisepolynomial.jl

@@ -25,6 +25,15 @@ function PiecewisePolynomialKernel(;v::Integer=0, maha::AbstractMatrix{<:Real})
     return PiecewisePolynomialKernel{v}(maha)
 end
 
+# Have to reconstruct the type parameter
+# See also https://github.com/FluxML/Functors.jl/issues/3#issuecomment-626747663
+function Functors.functor(::Type{<:PiecewisePolynomialKernel{V}}, x) where V
+    function reconstruct_kernel(xs)
+        return PiecewisePolynomialKernel{V}(xs.maha)
+    end
+    return (maha = x.maha,), reconstruct_kernel
+end
+
 _f(κ::PiecewisePolynomialKernel{0}, r, j) = 1
 _f(κ::PiecewisePolynomialKernel{1}, r, j) = 1 + (j + 1) * r
 _f(κ::PiecewisePolynomialKernel{2}, r, j) = 1 + (j + 2) * r + (j^2 + 4 * j + 3) / 3 * r.^2

src/basekernels/polynomial.jl

@@ -14,6 +14,8 @@ struct LinearKernel{Tc<:Real} <: SimpleKernel
     end
 end
 
+@functor LinearKernel
+
 kappa(κ::LinearKernel, xᵀy::Real) = xᵀy + first(κ.c)
 
 metric(::LinearKernel) = DotProduct()
@@ -38,6 +40,8 @@ struct PolynomialKernel{Td<:Real, Tc<:Real} <: SimpleKernel
     end
 end
 
+@functor PolynomialKernel
+
 kappa(κ::PolynomialKernel, xᵀy::Real) = (xᵀy + first(κ.c))^(first(κ.d))
 
 metric(::PolynomialKernel) = DotProduct()

src/basekernels/rationalquad.jl

@@ -15,6 +15,8 @@ struct RationalQuadraticKernel{Tα<:Real} <: SimpleKernel
     end
 end
 
+@functor RationalQuadraticKernel
+
 kappa(κ::RationalQuadraticKernel, d²::T) where {T<:Real} = (one(T)+d²/first(κ.α))^(-first(κ.α))
 metric(::RationalQuadraticKernel) = SqEuclidean()
 
@@ -38,6 +40,8 @@ struct GammaRationalQuadraticKernel{Tα<:Real, Tγ<:Real} <: SimpleKernel
     end
 end
 
+@functor GammaRationalQuadraticKernel
+
 kappa(κ::GammaRationalQuadraticKernel, d²::T) where {T<:Real} = (one(T)+d²^first(κ.γ)/first(κ.α))^(-first(κ.α))
 metric(::GammaRationalQuadraticKernel) = SqEuclidean()
 

src/kernels/kernelproduct.jl

@@ -39,6 +39,8 @@ function KernelProduct(kernel::Kernel, kernels::Kernel...)
     return KernelProduct((kernel, kernels...))
 end
 
+@functor KernelProduct
+
 Base.:*(k1::Kernel,k2::Kernel) = KernelProduct(k1, k2)
 
 function Base.:*(

src/kernels/kernelsum.jl

@@ -39,6 +39,8 @@ function KernelSum(kernel::Kernel, kernels::Kernel...)
     return KernelSum((kernel, kernels...))
 end
 
+@functor KernelSum
+
 Base.:+(k1::Kernel, k2::Kernel) = KernelSum(k1, k2)
 
 function Base.:+(

src/kernels/scaledkernel.jl

@@ -13,6 +13,8 @@ function ScaledKernel(kernel::Tk, σ²::Tσ²=1.0) where {Tk<:Kernel,Tσ²<:Real
     return ScaledKernel{Tk, Tσ²}(kernel, [σ²])
 end
 
+@functor ScaledKernel
+
 (k::ScaledKernel)(x, y) = first(k.σ²) * k.kernel(x, y)
 
 function kernelmatrix(κ::ScaledKernel, x::AbstractVector, y::AbstractVector)

src/kernels/tensorproduct.jl

@@ -20,6 +20,8 @@ function TensorProduct(kernel::Kernel, kernels::Kernel...)
     return TensorProduct((kernel, kernels...))
 end
 
+@functor TensorProduct
+
 Base.length(kernel::TensorProduct) = length(kernel.kernels)
 
 function (kernel::TensorProduct)(x, y)

src/kernels/transformedkernel.jl

@@ -9,6 +9,8 @@ struct TransformedKernel{Tk<:Kernel,Tr<:Transform} <: Kernel
     transform::Tr
 end
 
+@functor TransformedKernel
+
 (k::TransformedKernel)(x, y) = k.kernel(k.transform(x), k.transform(y))
 
 # Optimizations for scale transforms of simple kernels to save allocations:

src/trainable.jl

@@ -14,6 +14,8 @@ end
 
 ARDTransform(s::Real, dims::Integer) = ARDTransform(fill(s, dims))
 
+@functor ARDTransform
+
 function set!(t::ARDTransform{<:AbstractVector{T}}, ρ::AbstractVector{T}) where {T<:Real}
     @assert length(ρ) == dim(t) "Trying to set a vector of size $(length(ρ)) to ARDTransform of dimension $(dim(t))"
     t.v .= ρ

src/transform/ardtransform.jl

@@ -13,6 +13,8 @@ struct ChainTransform{V<:AbstractVector{<:Transform}} <: Transform
     transforms::V
 end
 
+@functor ChainTransform
+
 Base.length(t::ChainTransform) = length(t.transforms)
 
 # Constructor to create a chain transform with an array of parameters

src/transform/chaintransform.jl

@@ -13,6 +13,8 @@ struct FunctionTransform{F} <: Transform
     f::F
 end
 
+@functor FunctionTransform
+
 (t::FunctionTransform)(x) = t.f(x)
 
 _map(t::FunctionTransform, x::AbstractVector{<:Real}) = map(t.f, x)

src/transform/functiontransform.jl

@@ -17,6 +17,8 @@ struct LinearTransform{T<:AbstractMatrix{<:Real}} <: Transform
     A::T
 end
 
+@functor LinearTransform
+
 function set!(t::LinearTransform{<:AbstractMatrix{T}}, A::AbstractMatrix{T}) where {T<:Real}
     size(t.A) == size(A) ||
         error("size of the given matrix ", size(A), " and of the transformation matrix ",

src/transform/lineartransform.jl

@@ -15,6 +15,8 @@ function ScaleTransform(s::T=1.0) where {T<:Real}
     ScaleTransform{T}([s])
 end
 
+@functor ScaleTransform
+
 set!(t::ScaleTransform,ρ::Real) = t.s .= [ρ]
 
 (t::ScaleTransform)(x) = first(t.s) * x

src/transform/scaletransform.jl

@@ -26,6 +26,7 @@
         @test metric(ConstantKernel()) == KernelFunctions.Delta()
         @test metric(ConstantKernel(c=2.0)) == KernelFunctions.Delta()
         @test repr(k) == "Constant Kernel (c = $(c))"
+        test_params(k, ([c],))
         test_ADs(c->ConstantKernel(c=first(c)), [c])
     end
 end

test/basekernels/constant.jl

@@ -40,6 +40,7 @@
         @test KernelFunctions.iskroncompatible(k) == true
         test_ADs(γ -> GammaExponentialKernel(gamma=first(γ)), [γ], ADs = [:ForwardDiff, :ReverseDiff])
         @test_broken "Zygote gradient given γ"
+        test_params(k, ([γ],))
         #Coherence :
         @test GammaExponentialKernel(γ=1.0)(v1,v2) ≈ SqExponentialKernel()(v1,v2)
         @test GammaExponentialKernel(γ=0.5)(v1,v2) ≈ ExponentialKernel()(v1,v2)

test/basekernels/exponential.jl

@@ -24,4 +24,6 @@
     @test repr(k) == "Fractional Brownian Motion Kernel (h = $(h))"
     test_ADs(FBMKernel, ADs = [:ReverseDiff])
     @test_broken "Tests failing for kernelmatrix(k, x) for ForwardDiff and Zygote"
+
+    test_params(k, ([h],))
 end

test/basekernels/fbm.jl

@@ -13,4 +13,6 @@
     @test repr(k) == "Mahalanobis Kernel (size(P) = $(size(P)))"
     # test_ADs(P -> MahalanobisKernel(P=P), P)
     @test_broken "Nothing passes (problem with Mahalanobis distance in Distances)"
+
+    test_params(k, (P,))
 end

test/basekernels/maha.jl

@@ -16,6 +16,7 @@
         @test repr(k) == "Matern Kernel (ν = $(ν))"
         # test_ADs(x->MaternKernel(nu=first(x)),[ν])
         @test_broken "All fails (because of logabsgamma for ForwardDiff and ReverseDiff and because of nu for Zygote)"
+        test_params(k, ([ν],))
     end
     @testset "Matern32Kernel" begin
         k = Matern32Kernel()

test/basekernels/matern.jl

@@ -9,4 +9,5 @@
     @test repr(k) == "Periodic Kernel, length(r) = $(length(r)))"
     # test_ADs(r->PeriodicKernel(r =exp.(r)), log.(r), ADs = [:ForwardDiff, :ReverseDiff])
     @test_broken "Undefined adjoint for Sinus metric, and failing randomly for ForwardDiff and ReverseDiff"
+    test_params(k, (r,))
 end

test/basekernels/periodic.jl

@@ -34,4 +34,6 @@
     @test repr(k) == "Piecewise Polynomial Kernel (v = $(v), size(maha) = $(size(maha)))"
     # test_ADs(maha-> PiecewisePolynomialKernel(v=2, maha = maha), maha)
     @test_broken "Nothing passes (problem with Mahalanobis distance in Distances)"
+
+    test_params(k, (maha,))
 end