streamlining inits to WeightInitializers v1

Author: MartinuzziFrancesco

Project.toml

@@ -7,10 +7,8 @@ version = "0.10.4"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CellularAutomata = "878138dc-5b27-11ea-1a71-cb95d38d6b29"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
-Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 PartialFunctions = "570af359-4316-4cb7-8c74-252c00c2016b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -30,18 +28,16 @@ Aqua = "0.8"
 CellularAutomata = "0.0.2"
 DifferentialEquations = "7"
 Distances = "0.10"
-Distributions = "0.25.36"
 LIBSVM = "0.8"
 LinearAlgebra = "1.10"
 MLJLinearModels = "0.9.2, 0.10"
 NNlib = "0.8.4, 0.9"
-Optim = "1"
 PartialFunctions = "1.2"
 Random = "1.10"
 SafeTestsets = "0.1"
 Statistics = "1.10"
 Test = "1"
-WeightInitializers = "0.1.6"
+WeightInitializers = "1.0.4"
 julia = "1.10"
 
 [extras]
@@ -52,4 +48,4 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Test", "SafeTestsets", "Random", "DifferentialEquations", "MLJLinearModels", "LIBSVM"]
+test = ["Aqua", "Test", "SafeTestsets", "Random", "DifferentialEquations"]

src/ReservoirComputing.jl

@@ -10,7 +10,7 @@ using Optim
 using PartialFunctions
 using Random
 using Statistics
-using WeightInitializers
+using WeightInitializers: WeightInitializers, DeviceAgnostic
 
 export NLADefault, NLAT1, NLAT2, NLAT3
 export StandardStates, ExtendedStates, PaddedStates, PaddedExtendedStates
@@ -106,30 +106,35 @@ end
 
 __partial_apply(fn, inp) = fn$inp
 
-#fallbacks for initializers
+#fallbacks for initializers #eventually to remove once migrated to WeightInitializers.jl
 for initializer in (:rand_sparse, :delay_line, :delay_line_backward, :cycle_jumps,
     :simple_cycle, :pseudo_svd,
     :scaled_rand, :weighted_init, :informed_init, :minimal_init)
-    NType = ifelse(initializer === :rand_sparse, Real, Number)
-    @eval function ($initializer)(dims::Integer...; kwargs...)
-        return $initializer(WeightInitializers._default_rng(), Float32, dims...; kwargs...)
+    @eval begin
+        function ($initializer)(dims::Integer...; kwargs...)
+            return $initializer(Utils.default_rng(), Float32, dims...; kwargs...)
+        end
+        function ($initializer)(rng::AbstractRNG, dims::Integer...; kwargs...)
+            return $initializer(rng, Float32, dims...; kwargs...)
+        end
+        function ($initializer)(::Type{T}, dims::Integer...; kwargs...) where {T <: Number}
+            return $initializer(Utils.default_rng(), T, dims...; kwargs...)
+        end
+
+        # Partial application
+        function ($initializer)(rng::AbstractRNG; kwargs...)
+            return PartialFunction.Partial{Nothing}($initializer, rng, kwargs)
+        end
+        function ($initializer)(::Type{T}; kwargs...) where {T <: Number}
+            return PartialFunction.Partial{T}($initializer, nothing, kwargs)
+        end
+        function ($initializer)(rng::AbstractRNG, ::Type{T}; kwargs...) where {T <: Number}
+            return PartialFunction.Partial{T}($initializer, rng, kwargs)
+        end
+        function ($initializer)(; kwargs...)
+            return PartialFunction.Partial{Nothing}($initializer, nothing, kwargs)
+        end
     end
-    @eval function ($initializer)(rng::AbstractRNG, dims::Integer...; kwargs...)
-        return $initializer(rng, Float32, dims...; kwargs...)
-    end
-    @eval function ($initializer)(::Type{T},
-            dims::Integer...; kwargs...) where {T <: $NType}
-        return $initializer(WeightInitializers._default_rng(), T, dims...; kwargs...)
-    end
-    @eval function ($initializer)(rng::AbstractRNG; kwargs...)
-        return __partial_apply($initializer, (rng, (; kwargs...)))
-    end
-    @eval function ($initializer)(rng::AbstractRNG,
-            ::Type{T}; kwargs...) where {T <: $NType}
-        return __partial_apply($initializer, ((rng, T), (; kwargs...)))
-    end
-    @eval ($initializer)(; kwargs...) = __partial_apply(
-        $initializer, (; kwargs...))
 end
 
 #general

src/esn/esn_input_layers.jl

@@ -26,7 +26,7 @@ function scaled_rand(rng::AbstractRNG,
         dims::Integer...;
         scaling = T(0.1)) where {T <: Number}
     res_size, in_size = dims
-    layer_matrix = T.(rand(rng, Uniform(-scaling, scaling), res_size, in_size))
+    layer_matrix = T.(DeviceAgnostic.rand(rng, Uniform(-scaling, scaling), res_size, in_size))
     return layer_matrix
 end
 
@@ -65,11 +65,11 @@ function weighted_init(rng::AbstractRNG,
         scaling = T(0.1)) where {T <: Number}
     approx_res_size, in_size = dims
     res_size = Int(floor(approx_res_size / in_size) * in_size)
-    layer_matrix = zeros(T, res_size, in_size)
+    layer_matrix = DeviceAgnostic.zeros(T, res_size, in_size)
     q = floor(Int, res_size / in_size)
 
     for i in 1:in_size
-        layer_matrix[((i - 1) * q + 1):((i) * q), i] = rand(rng,
+        layer_matrix[((i - 1) * q + 1):((i) * q), i] = DeviceAgnostic.rand(rng,
             Uniform(-scaling, scaling),
             q)
     end
@@ -113,25 +113,25 @@ function informed_init(rng::AbstractRNG, ::Type{T}, dims::Integer...;
         throw(DimensionMismatch("in_size must be greater than model_in_size"))
     end
 
-    input_matrix = zeros(res_size, in_size)
-    zero_connections = zeros(in_size)
+    input_matrix = DeviceAgnostic.zeros(res_size, in_size)
+    zero_connections = DeviceAgnostic.zeros(in_size)
     num_for_state = floor(Int, res_size * gamma)
     num_for_model = floor(Int, res_size * (1 - gamma))
 
     for i in 1:num_for_state
         idxs = findall(Bool[zero_connections .== input_matrix[i, :]
                             for i in 1:size(input_matrix, 1)])
-        random_row_idx = idxs[rand(rng, 1:end)]
-        random_clm_idx = range(1, state_size, step = 1)[rand(rng, 1:end)]
-        input_matrix[random_row_idx, random_clm_idx] = rand(rng, Uniform(-scaling, scaling))
+        random_row_idx = idxs[DeviceAgnostic.rand(rng, 1:end)]
+        random_clm_idx = range(1, state_size, step = 1)[DeviceAgnostic.rand(rng, 1:end)]
+        input_matrix[random_row_idx, random_clm_idx] = DeviceAgnostic.rand(rng, Uniform(-scaling, scaling))
     end
 
     for i in 1:num_for_model
         idxs = findall(Bool[zero_connections .== input_matrix[i, :]
                             for i in 1:size(input_matrix, 1)])
-        random_row_idx = idxs[rand(rng, 1:end)]
-        random_clm_idx = range(state_size + 1, in_size, step = 1)[rand(rng, 1:end)]
-        input_matrix[random_row_idx, random_clm_idx] = rand(rng, Uniform(-scaling, scaling))
+        random_row_idx = idxs[DeviceAgnostic.rand(rng, 1:end)]
+        random_clm_idx = range(state_size + 1, in_size, step = 1)[DeviceAgnostic.rand(rng, 1:end)]
+        input_matrix[random_row_idx, random_clm_idx] = DeviceAgnostic.rand(rng, Uniform(-scaling, scaling))
     end
 
     return input_matrix
@@ -196,10 +196,10 @@ function _create_bernoulli(p::Number,
         weight::Number,
         rng::AbstractRNG,
         ::Type{T}) where {T <: Number}
-    input_matrix = zeros(T, res_size, in_size)
+    input_matrix = DeviceAgnostic.zeros(T, res_size, in_size)
     for i in 1:res_size
         for j in 1:in_size
-            rand(rng, Bernoulli(p)) ? (input_matrix[i, j] = weight) :
+            DeviceAgnostic.rand(rng, Bernoulli(p)) ? (input_matrix[i, j] = weight) :
             (input_matrix[i, j] = -weight)
         end
     end
@@ -216,16 +216,16 @@ function _create_irrational(irrational::Irrational,
     setprecision(BigFloat, Int(ceil(log2(10) * (res_size * in_size + start + 1))))
     ir_string = string(BigFloat(irrational)) |> collect
     deleteat!(ir_string, findall(x -> x == '.', ir_string))
-    ir_array = zeros(length(ir_string))
-    input_matrix = zeros(T, res_size, in_size)
+    ir_array = DeviceAgnostic.zeros(length(ir_string))
+    input_matrix = DeviceAgnostic.zeros(T, res_size, in_size)
 
     for i in 1:length(ir_string)
         ir_array[i] = parse(Int, ir_string[i])
     end
 
     for i in 1:res_size
         for j in 1:in_size
-            random_number = rand(rng, T)
+            random_number = DeviceAgnostic.rand(rng, T)
             input_matrix[i, j] = random_number < 0.5 ? -weight : weight
         end
     end

src/esn/esn_reservoirs.jl

@@ -66,7 +66,7 @@ function delay_line(rng::AbstractRNG,
         ::Type{T},
         dims::Integer...;
         weight = T(0.1)) where {T <: Number}
-    reservoir_matrix = zeros(T, dims...)
+    reservoir_matrix = DeviceAgnostic.zeros(T, dims...)
     @assert length(dims) == 2&&dims[1] == dims[2] "The dimensions must define a square matrix (e.g., (100, 100))"
 
     for i in 1:(dims[1] - 1)
@@ -107,7 +107,7 @@ function delay_line_backward(rng::AbstractRNG,
         weight = T(0.1),
         fb_weight = T(0.2)) where {T <: Number}
     res_size = first(dims)
-    reservoir_matrix = zeros(T, dims...)
+    reservoir_matrix = DeviceAgnostic.zeros(T, dims...)
 
     for i in 1:(res_size - 1)
         reservoir_matrix[i + 1, i] = weight
@@ -148,7 +148,7 @@ function cycle_jumps(rng::AbstractRNG,
         jump_weight::Number = T(0.1),
         jump_size::Int = 3) where {T <: Number}
     res_size = first(dims)
-    reservoir_matrix = zeros(T, dims...)
+    reservoir_matrix = DeviceAgnostic.zeros(T, dims...)
 
     for i in 1:(res_size - 1)
         reservoir_matrix[i + 1, i] = cycle_weight
@@ -194,7 +194,7 @@ function simple_cycle(rng::AbstractRNG,
         ::Type{T},
         dims::Integer...;
         weight = T(0.1)) where {T <: Number}
-    reservoir_matrix = zeros(T, dims...)
+    reservoir_matrix = DeviceAgnostic.zeros(T, dims...)
 
     for i in 1:(dims[1] - 1)
         reservoir_matrix[i + 1, i] = weight
@@ -246,9 +246,9 @@ function pseudo_svd(rng::AbstractRNG,
 
     while tmp_sparsity <= sparsity
         reservoir_matrix *= create_qmatrix(dims[1],
-            rand(1:dims[1]),
-            rand(1:dims[1]),
-            rand(T) * T(2) - T(1),
+            DeviceAgnostic.rand(1:dims[1]),
+            DeviceAgnostic.rand(1:dims[1]),
+            DeviceAgnostic.rand(T) * T(2) - T(1),
             T)
         tmp_sparsity = get_sparsity(reservoir_matrix, dims[1])
     end
@@ -258,17 +258,17 @@ end
 
 function create_diag(dim::Number, max_value::Number, ::Type{T};
         sorted::Bool = true, reverse_sort::Bool = false) where {T <: Number}
-    diagonal_matrix = zeros(T, dim, dim)
+    diagonal_matrix = DeviceAgnostic.zeros(T, dim, dim)
     if sorted == true
         if reverse_sort == true
-            diagonal_values = sort(rand(T, dim) .* max_value, rev = true)
+            diagonal_values = sort(DeviceAgnostic.rand(T, dim) .* max_value, rev = true)
             diagonal_values[1] = max_value
         else
-            diagonal_values = sort(rand(T, dim) .* max_value)
+            diagonal_values = sort(DeviceAgnostic.rand(T, dim) .* max_value)
             diagonal_values[end] = max_value
         end
     else
-        diagonal_values = rand(T, dim) .* max_value
+        diagonal_values = DeviceAgnostic.rand(T, dim) .* max_value
     end
 
     for i in 1:dim
@@ -283,7 +283,7 @@ function create_qmatrix(dim::Number,
         coord_j::Number,
         theta::Number,
         ::Type{T}) where {T <: Number}
-    qmatrix = zeros(T, dim, dim)
+    qmatrix = DeviceAgnostic.zeros(T, dim, dim)
 
     for i in 1:dim
         qmatrix[i, i] = 1.0