Moved trainable to trainable.jl and added Flux via Requires

theogf · theogf · commit f9d9bd0f9ec7 · 2020-03-11T18:02:25.000+01:00
diff --git a/src/KernelFunctions.jl b/src/KernelFunctions.jl
@@ -58,6 +58,7 @@ 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
diff --git a/src/generic.jl b/src/generic.jl
@@ -14,21 +14,7 @@ _scale(t::ScaleTransform, metric, x, y) = evaluate(metric, apply(t, x), apply(t,
 printshifted(io::IO,κ::Kernel,shift::Int) = print(io,"$κ")
 Base.show(io::IO,κ::Kernel) = print(io,nameof(typeof(κ)))
 
-function params(k::Kernel)
-    ps = []
-    params!(ps,k)
-    return ps
-end
-
-function params!(ps,k::Kernel)
-    for child in trainable(k)
-        params!(ps,k)
-    end
-end
-
-params!(ps,x::AbstractArray) = push!(ps,x)
-
-trainable(x) = ()
+_trainable(x) = ()
 
 ### Syntactic sugar for creating matrices and using kernel functions
 for k in subtypes(BaseKernel)
diff --git a/src/kernels/constant.jl b/src/kernels/constant.jl
@@ -41,8 +41,6 @@ struct ConstantKernel{Tc<:Real} <: BaseKernel
     end
 end
 
-trainable(k::ConstantKernel) = (k.c,)
-
 kappa(κ::ConstantKernel,x::Real) = first(κ.c)*one(x)
 
 metric(::ConstantKernel) = Delta()
diff --git a/src/kernels/exponential.jl b/src/kernels/exponential.jl
@@ -54,8 +54,6 @@ struct GammaExponentialKernel{Tγ<:Real} <: BaseKernel
     end
 end
 
-trainable(k::GammaExponentialKernel) = (γ,)
-
 kappa(κ::GammaExponentialKernel, d²::Real) = exp(-d²^first(κ.γ))
 iskroncompatible(::GammaExponentialKernel) = true
 metric(::GammaExponentialKernel) = SqEuclidean()
diff --git a/src/kernels/kernelproduct.jl b/src/kernels/kernelproduct.jl
@@ -14,9 +14,6 @@ struct KernelProduct <: Kernel
     kernels::Vector{Kernel}
 end
 
-params(k::KernelProduct) = params.(k.kernels)
-opt_params(k::KernelProduct) = opt_params.(k.kernels)
-
 Base.:*(k1::Kernel,k2::Kernel) = KernelProduct([k1,k2])
 Base.:*(k1::KernelProduct,k2::KernelProduct) = KernelProduct(vcat(k1.kernels,k2.kernels)) #TODO Add test
 Base.:*(k::Kernel,kp::KernelProduct) = KernelProduct(vcat(k,kp.kernels))
diff --git a/src/kernels/kernelsum.jl b/src/kernels/kernelsum.jl
@@ -25,9 +25,6 @@ function KernelSum(
     KernelSum(kernels, weights)
 end
 
-params(k::KernelSum) = (k.weights, params.(k.kernels))
-opt_params(k::KernelSum) = (k.weights, opt_params.(k.kernels))
-
 Base.:+(k1::Kernel, k2::Kernel) = KernelSum([k1, k2], weights = [1.0, 1.0])
 Base.:+(k1::ScaledKernel, k2::ScaledKernel) = KernelSum([kernel(k1), kernel(k2)], weights = [first(k1.σ), first(k2.σ)])
 Base.:+(k1::KernelSum, k2::KernelSum) =
diff --git a/src/kernels/matern.jl b/src/kernels/matern.jl
@@ -14,8 +14,6 @@ struct MaternKernel{Tν<:Real} <: BaseKernel
     end
 end
 
-trainable(k::MaternKernel) = (k.ν,)
-
 @inline function kappa(κ::MaternKernel, d::Real)
     ν = first(κ.ν)
     iszero(d) ? one(d) : exp((one(d)-ν)*logtwo-logabsgamma(ν)[1] + ν*log(sqrt(2ν)*d)+log(besselk(ν,sqrt(2ν)*d)))
diff --git a/src/kernels/polynomial.jl b/src/kernels/polynomial.jl
@@ -13,8 +13,6 @@ struct LinearKernel{Tc<:Real} <: BaseKernel
     end
 end
 
-trainable(k::LinearKernel) = (k.c,)
-
 kappa(κ::LinearKernel, xᵀy::Real) = xᵀy + first(κ.c)
 
 metric(::LinearKernel) = DotProduct()
@@ -36,8 +34,6 @@ struct PolynomialKernel{Td<:Real, Tc<:Real} <: BaseKernel
     end
 end
 
-trainable(k::PolynomialKernel) = (k.d,k.c)
-
 kappa(κ::PolynomialKernel, xᵀy::T) where {T<:Real} = (xᵀy + first(κ.c))^(first(κ.d))
 
 metric(::PolynomialKernel) = DotProduct()
diff --git a/src/kernels/rationalquad.jl b/src/kernels/rationalquad.jl
@@ -14,8 +14,6 @@ struct RationalQuadraticKernel{Tα<:Real} <: BaseKernel
     end
 end
 
-trainable(k::RationalQuadraticKernel) = (k.α,)
-
 kappa(κ::RationalQuadraticKernel, d²::T) where {T<:Real} = (one(T)+d²/first(κ.α))^(-first(κ.α))
 
 metric(::RationalQuadraticKernel) = SqEuclidean()
@@ -38,8 +36,6 @@ struct GammaRationalQuadraticKernel{Tα<:Real, Tγ<:Real} <: BaseKernel
     end
 end
 
-trainable(k::GammaRationalQuadraticKernel) = (k.α,k.γ)
-
 kappa(κ::GammaRationalQuadraticKernel, d²::T) where {T<:Real} = (one(T)+d²^first(κ.γ)/first(κ.α))^(-first(κ.α))
 
 metric(::GammaRationalQuadraticKernel) = SqEuclidean()
diff --git a/src/kernels/scaledkernel.jl b/src/kernels/scaledkernel.jl
@@ -12,8 +12,6 @@ kappa(k::ScaledKernel, x) = first(k.σ)*kappa(k.kernel, x)
 
 metric(k::ScaledKernel) = metric(k.kernel)
 
-trainable(k::ScaledKernel) = (k.σ,k.kernel)
-
 Base.:*(w::Real,k::Kernel) = ScaledKernel(k,w)
 
 Base.show(io::IO,κ::ScaledKernel) = printshifted(io,κ,0)
diff --git a/src/kernels/transformedkernel.jl b/src/kernels/transformedkernel.jl
@@ -27,8 +27,6 @@ kappa(κ::TransformedKernel, x) = kappa(κ.kernel, x)
 
 metric(κ::TransformedKernel) = metric(κ.kernel)
 
-trainable(κ::TransformedKernel) = (κ.transform,κ.kernel)
-
 Base.show(io::IO,κ::TransformedKernel) = printshifted(io,κ,0)
 
 function printshifted(io::IO,κ::TransformedKernel,shift::Int)
diff --git a/src/trainable.jl b/src/trainable.jl
@@ -0,0 +1,42 @@
+using .Flux: trainable
+
+Flux.trainable(::Kernel) = () # By default no parameters are returned
+Flux.trainable(::Transform) = ()
+
+### Base Kernels
+
+Flux.trainable(k::ConstantKernel) = (k.c,)
+
+Flux.trainable(k::GammaExponentialKernel) = (γ,)
+
+Flux.trainable(k::GammaRationalQuadraticKernel) = (k.α, k.γ)
+
+Flux.trainable(k::MaternKernel) = (k.ν,)
+
+Flux.trainable(k::LinearKernel) = (k.c,)
+
+Flux.trainable(k::PolynomialKernel) = (k.d, k.c)
+
+Flux.trainable(k::RationalQuadraticKernel) = (k.α,)
+
+#### Composite kernels
+
+Flux.trainable(κ::KernelProduct) = k.kernels
+
+Flux.trainable(κ::KernelSum) = (k.weights, k.kernels) #To check
+
+Flux.trainable(k::ScaledKernel) = (k.σ, k.kernel)
+
+Flux.trainable(κ::TransformedKernel) = (κ.transform, κ.kernel)
+
+### Transforms
+
+Flux.trainable(t::ARDTransform) = (t.v,)
+
+Flux.trainable(t::ChainTransform) = t.transforms
+
+Flux.trainable(t::FunctionTransform) = (t.f,)
+
+Flux.trainable(t::LowRankTransform) = (t.proj,)
+
+Flux.trainable(t::ScaleTransform) = (t.s,)
diff --git a/src/transform/ardtransform.jl b/src/transform/ardtransform.jl
@@ -25,7 +25,6 @@ function set!(t::ARDTransform{T},ρ::AbstractVector{T}) where {T<:Real}
     t.v .= ρ
 end
 
-trainable(t::ARDTransform) = (t.v,)
 dim(t::ARDTransform) = length(t.v)
 
 function apply(t::ARDTransform,X::AbstractMatrix{<:Real};obsdim::Int = defaultobs)
diff --git a/src/transform/chaintransform.jl b/src/transform/chaintransform.jl
@@ -32,10 +32,8 @@ function apply(t::ChainTransform,X::T;obsdim::Int=defaultobs) where {T}
 end
 
 set!(t::ChainTransform,θ) = set!.(t.transforms,θ)
-trainable(t::ChainTransform) = t.transforms
 duplicate(t::ChainTransform,θ) = ChainTransform(duplicate.(t.transforms,θ))
 
-
 Base.:∘(t₁::Transform,t₂::Transform) = ChainTransform([t₂,t₁])
 Base.:∘(t::Transform,tc::ChainTransform) = ChainTransform(vcat(tc.transforms,t)) #TODO add test
 Base.:∘(tc::ChainTransform,t::Transform) = ChainTransform(vcat(t,tc.transforms))
diff --git a/src/transform/functiontransform.jl b/src/transform/functiontransform.jl
@@ -15,4 +15,3 @@ end
 apply(t::FunctionTransform, X::T; obsdim::Int = defaultobs) where {T} = mapslices(t.f, X, dims = feature_dim(obsdim))
 
 duplicate(t::FunctionTransform,f) = FunctionTransform(f)
-trainable(t::FunctionTransform) = (t.f,)
diff --git a/src/transform/lowranktransform.jl b/src/transform/lowranktransform.jl
@@ -16,7 +16,6 @@ function set!(t::LowRankTransform{<:AbstractMatrix{T}},M::AbstractMatrix{T}) whe
     t.proj .= M
 end
 
-trainable(t::LowRankTransform) = (t.proj,)
 
 Base.size(tr::LowRankTransform,i::Int) = size(tr.proj,i)
 Base.size(tr::LowRankTransform) = size(tr.proj) #  TODO Add test
diff --git a/src/transform/scaletransform.jl b/src/transform/scaletransform.jl
@@ -16,7 +16,7 @@ function ScaleTransform(s::T=1.0) where {T<:Real}
 end
 
 set!(t::ScaleTransform,ρ::Real) = t.s .= [ρ]
-trainable(t::ScaleTransform) = (t.s,)
+_trainable(t::ScaleTransform) = (t.s,)
 dim(str::ScaleTransform) = 1
 
 apply(t::ScaleTransform,x::AbstractVecOrMat;obsdim::Int=defaultobs) = first(t.s) * x
diff --git a/src/transform/selecttransform.jl b/src/transform/selecttransform.jl
@@ -24,7 +24,6 @@ set!(t::SelectTransform{<:AbstractVector{T}},dims::AbstractVector{T}) where {T<:
 
 duplicate(t::SelectTransform,θ) = t
 
-
 Base.maximum(t::SelectTransform) = maximum(t.select)
 
 function apply(t::SelectTransform, X::AbstractMatrix{<:Real}; obsdim::Int = defaultobs)
diff --git a/test/test_transform.jl b/test/test_transform.jl
@@ -25,6 +25,7 @@ f(x) = sin.(x)
         s2 = 2.0
         KernelFunctions.set!(t,s2)
         @test all(t.s.==[s2])
+        @test isequal(ScaleTransform(s),ScaleTransform(s))
     end
     ## Test ARD Transform
     @testset "ARDTransform" begin
@@ -35,17 +36,20 @@ f(x) = sin.(x)
         newv1 = rand(5)
         KernelFunctions.set!(vt1,newv1)
         @test all(vt1.v .== newv1)
+        @test ARDTransform(s,dims[2]).v == ARDTransform(s*ones(dims[2])).v
         @test_throws DimensionMismatch KernelFunctions.apply(vt1,rand(3,4))
     end
     ## Test LowRankTransform
     @testset "LowRankTransform" begin
         tp = LowRankTransform(P)
         @test all(KernelFunctions.apply(tp,X,obsdim=2).==P*X)
         @test all(KernelFunctions.apply(tp,x).==P*x)
-        @test all(KernelFunctions.params(tp).==P)
+        @test tp.proj == P
         P2 = rand(5,10)
         KernelFunctions.set!(tp,P2)
         @test all(tp.proj.==P2)
+        @test_throws AssertionError KernelFunctions.set!(tp,rand(6,10))
+        @test_throws DimensionMismatch KernelFunctions.apply(tp,rand(11,3))
     end
     ## Test FunctionTransform
     @testset "FunctionTransform" begin
