Added trainable for all transforms

Author: theogf

src/transform/ardtransform.jl

@@ -25,7 +25,7 @@ function set!(t::ARDTransform{T},ρ::AbstractVector{T}) where {T<:Real}
     t.v .= ρ
 end
 
-params(t::ARDTransform) = t.v
+trainable(t::ARDTransform) = (t.v,)
 dim(t::ARDTransform) = length(t.v)
 
 function apply(t::ARDTransform,X::AbstractMatrix{<:Real};obsdim::Int = defaultobs)

src/transform/chaintransform.jl

@@ -32,7 +32,7 @@ function apply(t::ChainTransform,X::T;obsdim::Int=defaultobs) where {T}
 end
 
 set!(t::ChainTransform,θ) = set!.(t.transforms,θ)
-params(t::ChainTransform) = (params.(t.transforms))
+trainable(t::ChainTransform) = t.transforms
 duplicate(t::ChainTransform,θ) = ChainTransform(duplicate.(t.transforms,θ))
 
 

src/transform/functiontransform.jl

@@ -4,7 +4,7 @@ FunctionTransform
     f(x) = abs.(x)
     tr = FunctionTransform(f)
 ```
-Take a function `f` as an argument which is going to act on each vector individually.
+Take a function or object `f` as an argument which is going to act on each vector individually.
 Make sure that `f` is supposed to act on a vector by eventually using broadcasting
 For example `f(x)=sin(x)` -> `f(x)=sin.(x)`
 """
@@ -15,4 +15,4 @@ 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)
-params(t::FunctionTransform) = t.f
+trainable(t::FunctionTransform) = (t.f,)

src/transform/lowranktransform.jl

@@ -16,7 +16,7 @@ function set!(t::LowRankTransform{<:AbstractMatrix{T}},M::AbstractMatrix{T}) whe
     t.proj .= M
 end
 
-params(t::LowRankTransform) = t.proj
+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

src/transform/scaletransform.jl

@@ -16,7 +16,7 @@ function ScaleTransform(s::T=1.0) where {T<:Real}
 end
 
 set!(t::ScaleTransform,ρ::Real) = t.s .= [ρ]
-params(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

src/transform/selecttransform.jl

@@ -22,8 +22,6 @@ end
 
 set!(t::SelectTransform{<:AbstractVector{T}},dims::AbstractVector{T}) where {T<:Int} = t.select .= dims
 
-params(t::SelectTransform) = t.select
-
 duplicate(t::SelectTransform,θ) = t
 
 

src/transform/transform.jl

@@ -7,15 +7,19 @@ include("functiontransform.jl")
 include("selecttransform.jl")
 include("chaintransform.jl")
 
+"""
+`apply(t::Transform, x; obsdim::Int=defaultobs)`
+Apply the transform `t` per slice on the array `x`
+"""
+apply
+
 """
 IdentityTransform
 Return exactly the input
 """
 struct IdentityTransform <: Transform end
 
-params(t::IdentityTransform) = nothing
-
-apply(t::IdentityTransform, x; obsdim::Int=defaultobs) = x #TODO add test
+apply(t::IdentityTransform, x; obsdim::Int=defaultobs) = x
 
 ### TODO Maybe defining adjoints could help but so far it's not working
 
@@ -32,9 +36,9 @@ apply(t::IdentityTransform, x; obsdim::Int=defaultobs) = x #TODO add test
 
 # @adjoint transform(t::ScaleTransform{<:AbstractVector{<:Real}},x::AbstractVector{<:Real}) = transform(t,x),Δ->(ScaleTransform(nothing),t.s.*Δ)
 #
-#     @adjoint transform(t::ScaleTransform{<:AbstractVector{<:Real}},X::AbstractMatrix{<:Real},obsdim::Int) = transform(t,X,obsdim),Δ->begin
+#     @adjoint transform(t::ARDTransform{<:Real},X::AbstractMatrix{<:Real},obsdim::Int) = transform(t,X,obsdim),Δ->begin
 #     @show Δ,size(Δ);
-#     return (obsdim == 1 ? ScaleTransform()Δ'.*X : ScaleTransform()Δ.*X,transform(t,Δ,obsdim),nothing)
+#     return (obsdim == 1 ? ARD()Δ'.*X : ScaleTransform()Δ.*X,transform(t,Δ,obsdim),nothing)
 #     end
 #
 # @adjoint transform(t::ScaleTransform{T},x::AbstractVecOrMat,obsdim::Int) where {T<:Real} = transform(t,x), Δ->(ScaleTransform(one(T)),t.s.*Δ,nothing)