Merge pull request #586 from MilkshakeForReal/patch-1

Don't initialize parameters twice

Author: ChrisRackauckas

src/discretize.jl

@@ -437,8 +437,8 @@ function SciMLBase.symbolic_discretize(pde_system::PDESystem,
                 end
             else
                 x = map(chain) do x
-                    _x = ComponentArrays.ComponentArray(Lux.setup(Random.default_rng(),
-                                                                  x)[1])
+                    _x = ComponentArrays.ComponentArray(Lux.initialparameters(Random.default_rng(),
+                                                                              x))
                     Float64.(_x) # No ComponentArray GPU support
                 end
                 names = ntuple(i -> depvars[i], length(chain))
@@ -451,8 +451,8 @@ function SciMLBase.symbolic_discretize(pde_system::PDESystem,
                 init_params = init_params isa Array ? Float64.(init_params) :
                               init_params
             else
-                init_params = Float64.(ComponentArrays.ComponentArray(Lux.setup(Random.default_rng(),
-                                                                                chain)[1]))
+                init_params = Float64.(ComponentArrays.ComponentArray(Lux.initialparameters(Random.default_rng(),
+                                                                                            chain)))
             end
         end
     else

src/pinn_types.jl

@@ -346,7 +346,7 @@ mutable struct Phi{C, S}
     f::C
     st::S
     function Phi(chain::Lux.AbstractExplicitLayer)
-        ps, st = Lux.setup(Random.default_rng(), chain)
+        st = Lux.initialstates(Random.default_rng(), chain)
         new{typeof(chain), typeof(st)}(chain, st)
     end
     function Phi(chain::Flux.Chain)