Merge pull request #116 from PumasAI/initparamsbangcorrectargorder

fix arg order of `init_params!`

Author: chriselrod

Project.toml

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

examples/mnist_lenet.jl

@@ -48,19 +48,19 @@ G = SimpleChains.alloc_threaded_grad(lenetloss);
 SimpleChains.accuracy_and_loss(lenetloss, xtrain, p),
 SimpleChains.accuracy_and_loss(lenetloss, xtest, ytest, p)
 
-# SimpleChains.init_params!(lenet, p);
+# SimpleChains.init_params!(p, lenet);
 @time SimpleChains.train_batched!(G, p, lenetloss, xtrain, SimpleChains.ADAM(3e-4), 10);
 SimpleChains.accuracy_and_loss(lenetloss, xtrain, p),
 SimpleChains.accuracy_and_loss(lenetloss, xtest, ytest, p)
 
 
 
 # lenet.memory .= 0;
-SimpleChains.init_params!(lenet, p);
+SimpleChains.init_params!(p, lenet);
 @time SimpleChains.train_batched!(G, p, lenetloss, xtrain, SimpleChains.ADAM(3e-4), 10);
 SimpleChains.accuracy_and_loss(lenetloss, xtrain, p),
 SimpleChains.accuracy_and_loss(lenetloss, xtest, ytest, p)
-SimpleChains.init_params!(lenet, p);
+SimpleChains.init_params!(p, lenet);
 @time SimpleChains.train_batched!(G, p, lenetloss, xtrain, SimpleChains.ADAM(3e-4), 10);
 SimpleChains.accuracy_and_loss(lenetloss, xtrain, p),
 SimpleChains.accuracy_and_loss(lenetloss, xtest, ytest, p)
@@ -75,11 +75,11 @@ lenetloss.memory .= 0x00;
 @time valgrad!(g1, lenetloss, xtrain, p)
 g0 == g1
 lenet.memory .= 0;
-SimpleChains.init_params!(lenet, p);
+SimpleChains.init_params!(p, lenet);
 @time SimpleChains.train_batched!(G, p, lenetloss, xtrain, SimpleChains.ADAM(3e-4), 10);
 SimpleChains.accuracy_and_loss(lenetloss, xtrain, p),
 SimpleChains.accuracy_and_loss(lenetloss, xtest, ytest, p)
-SimpleChains.init_params!(lenet, p);
+SimpleChains.init_params!(p, lenet);
 @time SimpleChains.train_batched!(G, p, lenetloss, xtrain, SimpleChains.ADAM(3e-4), 10);
 SimpleChains.accuracy_and_loss(lenetloss, xtrain, p),
 SimpleChains.accuracy_and_loss(lenetloss, xtest, ytest, p)

src/penalty.jl

@@ -47,8 +47,8 @@ end
 function init_params(Λ::AbstractPenalty, ::Type{T}; rng::AbstractRNG=local_rng()) where {T}
   init_params(getchain(Λ), nothing, T; rng)
 end
-function init_params!(Λ::AbstractPenalty, x, id = nothing; rng::AbstractRNG=local_rng())
-  init_params!(getchain(Λ), x, id; rng)
+function init_params!(x, Λ::AbstractPenalty, id = nothing; rng::AbstractRNG=local_rng())
+  init_params!(x, getchain(Λ), id; rng)
 end
 
 target(c::AbstractPenalty) = target(getchain(c))

src/simple_chain.jl

@@ -1,4 +1,5 @@
 
+
 struct InputDimUnknown end
 const InputDim = Union{InputDimUnknown,Tuple{Vararg{Integer}}}
 
@@ -387,8 +388,8 @@ end
 Randomly initializes parameter vector `p` with input dim `id`. Input dim does not need to be specified if these were provided to the chain object itself.
 See the documentation of the individual layers to see how they are initialized, but it is generally via (Xavier) Glorot uniform or normal distributions.
 """
-function init_params!(
-  chn::SimpleChain, x::AbstractVector, id = nothing; rng::AbstractRNG = local_rng()
+@inline function init_params!(
+  x::AbstractVector, chn::SimpleChain, id = nothing; rng::AbstractRNG = local_rng()
 )
   GC.@preserve x _init_params!(chn.layers, pointer(x), chain_input_dims(chn, id), rng)
   return x
@@ -398,14 +399,14 @@ function _init_params!(layers::Tuple, p::Ptr, id, rng::AbstractRNG)
   _init_params!(Base.tail(layers), p, od, rng)
 end
 _init_params!(::Tuple{}, p::Ptr, _, ::AbstractRNG) = nothing
-function init_params(
+@inline function init_params(
   Λ::SimpleChain,
   id::Union{Nothing,InputDim} = nothing,
   ::Type{T} = Float32;
   rng::AbstractRNG=local_rng()
 ) where {T}
   _id = chain_input_dims(Λ, id)
-  init_params!(Λ, StrideArray{T}(undef, numparam(Λ, id)), chain_input_dims(Λ, _id); rng)
+  init_params!(StrideArray{T}(undef, numparam(Λ, id)), Λ, chain_input_dims(Λ, _id); rng)
 end
 
 """

test/runtests.jl

@@ -235,12 +235,8 @@ InteractiveUtils.versioninfo(verbose=true)
         @show isapprox(g, gfdd, rtol = 1e-8)
       end
       # let g=g, sc=sc, x=x, p=p
-      @test iszero(
-        countallocations!(g, FrontLastPenalty(sc, L2Penalty(2.3), NoPenalty()), x, p),
-      )
-      @test iszero(
-        countallocations!(g, FrontLastPenalty(scd, L2Penalty(2.3), L1Penalty(0.45)), x, p),
-      )
+      @test countallocations!(g, FrontLastPenalty(sc, L2Penalty(2.3), NoPenalty()), x, p) == 0
+      @test countallocations!(g, FrontLastPenalty(scd, L2Penalty(2.3), L1Penalty(0.45)), x, p) == 0
       # @test iszero(@allocated(valgrad!(g, sc, x, p)))
 
       td = TurboDense{true}(tanh, static(8))