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)