Merge pull request #85 from PumasAI/biasandweightviews

Add bias and weight views

Project.toml

@@ -1,7 +1,7 @@
 name = "SimpleChains"
 uuid = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
 authors = ["Chris Elrod <[email protected]> and contributors"]
-version = "0.2.8"
+version = "0.2.9"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/simple_chain.jl

@@ -31,27 +31,6 @@ struct SimpleChain{N,I<:InputDim,L<:Tuple{Vararg{Any,N}}}
   memory::Vector{UInt8}
 end
 
-#=
-input_dims(_) = nothing
-function _check_input_dims(x, i)
-  d = input_dims(x)
-  d === nothing || d == i || throw(ArgumentError("Input size of one layer did not match the next."))
-end
-function _input_dims(t::Tuple{L,Vararg}) where {L}
-  l = first(t)
-  d = input_dims(l)
-  d === nothing ? _input_dims(Base.tail(t)) : d
-end 
-
-
-_verify_chain(::Tuple{}, _) = nothing
-function _verify_chain(layers::Tuple{L,Vararg}, inputdim = _input_dims(layers)) where {L}
-  l = first(layers)
-  _check_input_dims(l, inputdim)
-  d = output_size(Val(Float32), l, (inputdim,))[2][1]
-  _verify_chain(Base.tail(layers), d)
-end
-=#
 chain_input_dims(c::SimpleChain) = c.inputdim
 
 SimpleChain(input_dim::Integer, l::Vararg) = SimpleChain((input_dim,), l, UInt8[])

src/utils.jl

@@ -233,15 +233,56 @@ Returns a tuple of the parameters of the SimpleChain `sc`, as a view of the para
 """
 function params(sc::SimpleChain, p::AbstractVector, inputdim = nothing)
   @unpack layers = sc
-  A = _params(layers, pointer(p), chain_input_dims(sc, inputdim))
+  A = _walk_chain(Val{:param}(), layers, pointer(p), chain_input_dims(sc, inputdim))
   _add_memory(A, p)
 end
-_params(::Tuple{}, _, __) = ()
-function _params(layers, p, inputdim)
-  A, p, outputdim = _getparams(first(layers), p, inputdim)
-  B = _params(Base.tail(layers), p, outputdim)
+"""
+    weights(sc::SimpleChain, p::AbstractVector, inputdim = nothing)
+
+Returns a tuple of the weights (parameters other than biases) of the SimpleChain `sc`, as a view of the parameter vector `p`.
+"""
+function weights(sc::SimpleChain, p::AbstractVector, inputdim = nothing)
+  @unpack layers = sc
+  A = _walk_chain(Val{:weight}(), layers, pointer(p), chain_input_dims(sc, inputdim))
+  _add_memory(A, p)
+end
+"""
+    biases(sc::SimpleChain, p::AbstractVector, inputdim = nothing)
+
+Returns a tuple of the biases of the SimpleChain `sc`, as a view of the parameter vector `p`.
+"""
+function biases(sc::SimpleChain, p::AbstractVector, inputdim = nothing)
+  @unpack layers = sc
+  A = _walk_chain(Val{:bias}(), layers, pointer(p), chain_input_dims(sc, inputdim))
+  _add_memory(A, p)
+end
+
+# definitions that happen to be right in most cases to save up
+# from implementing too much
+_get(::Val{:param}, _, x) = x
+_get(::Val{:weight}, _, ::Nothing) = nothing
+_get(::Val{:weight}, _, x) = x
+_get(::Val{:weight}, _, x::Tuple{A,B}) where {A,B} = first(x)
+_get(::Val{:bias}, _, ::Nothing) = nothing
+_get(::Val{:bias}, _, x) = nothing
+_get(::Val{:bias}, _, x::Tuple{A,B}) where {A,B} = last(x)
+@inline function _getparams(f::F, layer, p, inputdim) where {F}
+  A, p, outputdim = _getparams(layer, p, inputdim)
+  _get(f, layer, A), p, outputdim
+end
+# TODO: support nesting simple chains; below definition should enable recursive params
+#=
+function _getparams(f::F, layer::Union{AbstractPenalty,SimpleChain}, p, inputdim) where {F}
+  _walk_chain(f, layer, p, inputdim)
+end
+=#
+_walk_chain(___, ::Tuple{}, _, __) = ()
+function _walk_chain(f::F, layers, p, inputdim) where {F}
+  A, p, outputdim = _getparams(f, first(layers), p, inputdim)
+  B = _walk_chain(f, Base.tail(layers), p, outputdim)
   (A, B...)
 end
+
 _add_memory(A::PtrArray, p) = StrideArray(A, p)
 _add_memory(::Tuple{}, _) = ()
 function _add_memory(t::Tuple, p)

test/runtests.jl

@@ -492,11 +492,17 @@ InteractiveUtils.versioninfo(verbose=true)
       ),
     )
     p = SimpleChains.init_params(sc);
-    n0, (W1, b1), n2, W3 = SimpleChains.params(sc,p)
+    n0, (W1, b1), n2, W3 = SimpleChains.params(sc,p);
     @test n0 === n2 === nothing
     @test W1 == reshape(view(p,1:24*8),(8,24))
     @test b1 == view(p,24*8+1:25*8)
     @test W3 == reshape(@view(p[25*8+1:end]),(2,8))
+    n01, W11, n21, W31 = SimpleChains.weights(sc,p);
+    n02, b12, n22, n3 = SimpleChains.biases(sc,p);
+    @test n01 === n21 === n02 === n22 === n3
+    @test W11 === W1
+    @test W31 === W3
+    @test b12 === b1
   end
 end
 # TODO: test ambiguities once ForwardDiff fixes them, or once ForwardDiff is dropped