Merge pull request #1499 from DevJac/kwargs_lambda_fix

Fix layer init functions kwargs getting overwritten

Author: DhairyaLGandhi

src/utils.jl

@@ -135,7 +135,7 @@ function kaiming_uniform(rng::AbstractRNG, dims...; gain = √2)
 end
 
 kaiming_uniform(dims...; kwargs...) = kaiming_uniform(Random.GLOBAL_RNG, dims...; kwargs...)
-kaiming_uniform(rng::AbstractRNG; kwargs...) = (dims...; kwargs...) -> kaiming_uniform(rng, dims...; kwargs...)
+kaiming_uniform(rng::AbstractRNG; init_kwargs...) = (dims...; kwargs...) -> kaiming_uniform(rng, dims...; init_kwargs..., kwargs...)
 
 """
     kaiming_normal([rng=GLOBAL_RNG], dims...; gain = √2)
@@ -172,7 +172,7 @@ function kaiming_normal(rng::AbstractRNG, dims...; gain = √2f0)
 end
 
 kaiming_normal(dims...; kwargs...) = kaiming_normal(Random.GLOBAL_RNG, dims...; kwargs...)
-kaiming_normal(rng::AbstractRNG; kwargs...) = (dims...; kwargs...) -> kaiming_normal(rng, dims...; kwargs...)
+kaiming_normal(rng::AbstractRNG; init_kwargs...) = (dims...; kwargs...) -> kaiming_normal(rng, dims...; init_kwargs..., kwargs...)
 
 """
     sparse_init([rng=GLOBAL_RNG], dims...; sparsity, std = 0.01)
@@ -216,7 +216,7 @@ function sparse_init(rng::AbstractRNG, dims...; sparsity, std = 0.01)
 end
 
 sparse_init(dims...; kwargs...) = sparse_init(Random.GLOBAL_RNG, dims...; kwargs...)
-sparse_init(rng::AbstractRNG; kwargs...) = (dims...; kwargs...) -> sparse_init(rng, dims...; kwargs...)
+sparse_init(rng::AbstractRNG; init_kwargs...) = (dims...; kwargs...) -> sparse_init(rng, dims...; init_kwargs..., kwargs...)
 
 ones(T::Type, dims...) = Base.ones(T, dims...)
 zeros(T::Type, dims...) = Base.zeros(T, dims...)

test/utils.jl

@@ -119,6 +119,22 @@ end
       @test eltype(v) == Float32
     end
   end
+
+  @testset "partial_application" begin
+    big = 1e9
+
+    partial_ku = kaiming_uniform(gain=big)
+    @test maximum(partial_ku(8, 8)) > big / 2
+    @test maximum(partial_ku(8, 8, gain=1)) < big / 2
+
+    partial_kn = kaiming_normal(gain=big)
+    @test maximum(partial_kn(8, 8)) > big / 2
+    @test maximum(partial_kn(8, 8, gain=1)) < big / 2
+
+    partial_si = sparse_init(sparsity=1)
+    @test maximum(partial_si(8, 8)) == 0
+    @test maximum(partial_si(8, 8, sparsity=0)) > 0
+  end
 end
 
 @testset "Params" begin