replace usage of deprecated Flux.params to Flux.trainables

Author: Giovanni3A

src/flux_utils.jl

@@ -7,7 +7,7 @@ Extract the parameters of a Flux model (Flux.Chain or Flux.Dense) into a single
 vector.
 """
 function extract_flux_params(model::Union{Flux.Chain,Flux.Dense})
-    θ = Flux.params(model)
+    θ = Flux.trainables(model)
     return reduce(vcat, [vec(p) for p in θ])
 end
 
@@ -21,7 +21,7 @@ function fix_flux_params_single_model(
     θ::Vector{<:Real},
 )
     i = 1
-    for p in Flux.params(model)
+    for p in Flux.trainables(model)
         psize = prod(size(p))
         p .= reshape(θ[i:i+psize-1], size(p))
         i += psize
@@ -38,7 +38,7 @@ of parameters.
 function fix_flux_params_multi_model(models, θ::Vector{<:Real})
     i = 1
     for model in models
-        for p in Flux.params(model)
+        for p in Flux.trainables(model)
             psize = prod(size(p))
             p .= reshape(θ[i:i+psize-1], size(p))
             i += psize
@@ -54,8 +54,8 @@ Check if a Flux layer has parameters.
 """
 function has_params(layer)
     try
-        # Attempt to get parameters; if it works and isn't empty, return true
-        return !isempty(Flux.params(layer))
+        # Attempt to get trainable parameters; if it works and isn't empty, return true
+        return !isempty(Flux.trainable(layer))
     catch e
         # If there is an error (e.g. method not matching), assume no parameters
         return false

src/optimizers/bilevel.jl

@@ -177,10 +177,10 @@ function solve_bilevel(
     ilayer = 1
     for layer in model.forecast.networks[1]
         if has_params(layer)
-            for p in Flux.params(layer.weight)
+            for p in Flux.trainables(layer.weight)
                 p .= value.(predictive_model_vars[ilayer][:W])
             end
-            for p in Flux.params(layer.bias)
+            for p in Flux.trainables(layer.bias)
                 p .= value.(predictive_model_vars[ilayer][:b])
             end
         end

src/predictive_model.jl

@@ -221,5 +221,5 @@ function apply_gradient!(
 )
     loss3(m, X) = mean(dCdy'm(X'))
     grad = Zygote.gradient(loss3, model, X)[1]
-    Optimisers.update!(opt_state, model, grad)
+    return Optimisers.update!(opt_state, model, grad)
 end

test/test_predictive_model.jl

@@ -26,9 +26,9 @@ out_size = 2
         ones((1, in_size)),
         Flux.setup(Flux.Descent(0.1), forecaster),
     )
-    @test Flux.params(forecaster.networks[1])[1] ==
+    @test Flux.trainables(forecaster)[1] ==
           0.9 * ones((out_size, in_size))
-    @test Flux.params(forecaster.networks[1])[2] == 0.9 * ones(out_size)
+    @test Flux.trainables(forecaster)[2] == 0.9 * ones(out_size)
 end
 
 @testset "Single-Chain" begin
@@ -55,9 +55,9 @@ end
         ones((1, in_size)),
         Flux.setup(Flux.Descent(0.1), forecaster),
     )
-    @test Flux.params(forecaster.networks[1])[1] ==
+    @test Flux.trainables(forecaster)[1] ==
           0.9 * ones((out_size, in_size))
-    @test Flux.params(forecaster.networks[1])[2] == 0.9 * ones(out_size)
+    @test Flux.trainables(forecaster)[2] == 0.9 * ones(out_size)
 end
 
 @testset "Multi-Variate-Dense" begin
@@ -86,9 +86,9 @@ end
         ones((1, in_size)),
         Flux.setup(Flux.Descent(0.1), forecaster),
     )
-    @test Flux.params(forecaster.networks[1])[1] ==
+    @test Flux.trainables(forecaster)[1] ==
           0.8 * ones((model_out_size, model_in_size))
-    @test Flux.params(forecaster.networks[1])[2] == 0.8 * ones(model_out_size)
+    @test Flux.trainables(forecaster)[2] == 0.8 * ones(model_out_size)
 end
 
 @testset "Multi-Model-Dense" begin
@@ -124,10 +124,10 @@ end
         ones((1, in_size)),
         Flux.setup(Flux.Descent(0.1), forecaster),
     )
-    @test Flux.params(forecaster.networks[1])[1] ==
+    @test Flux.trainables(forecaster)[1] ==
           0.9 * ones((model_out_size, model_in_size))
-    @test Flux.params(forecaster.networks[1])[2] == 0.9 * ones(model_out_size)
-    @test Flux.params(forecaster.networks[2])[1] ==
+    @test Flux.trainables(forecaster)[2] == 0.9 * ones(model_out_size)
+    @test Flux.trainables(forecaster)[3] ==
           0.9 * ones((model_out_size, model_in_size))
-    @test Flux.params(forecaster.networks[2])[2] == 0.9 * ones(model_out_size)
+    @test Flux.trainables(forecaster)[4] == 0.9 * ones(model_out_size)
 end