Merge pull request FluxML#1 from FluxML/master

Rebase
Vonewman · Feb 8, 2021 · 5e801e2 · 5e801e2
2 parents 418b316 + cdb445c
commit 5e801e2
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/src/utils.jl b/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...)
 
 """
     orthogonal([rng=GLOBAL_RNG], dims...; gain = 1)
@@ -275,7 +275,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...)

diff --git a/test/utils.jl b/test/utils.jl
@@ -134,6 +134,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