Freezing neural network parameters in Universal Differential Equation framework

Specific Domains Modelling & Simulations
1

Hi,
I’m working on hybrid systems with the SciML framework, approximating one of the derivatives of an ODE system with a neural network. When I attempt, using Lux.freeze, to freeze some of the network parameters that I don’t want to be optimized (for example the weights in the snippet that I’ve attached), I get the following error during optimization:

ERROR: MethodError: no method matching _merge(::ReverseDiff.TrackedArray{Float64, Float64, 1, ComponentVector{Float64, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{Axis{(weight = ViewAxis(1:10, ShapedA …

snippet_fail_frozen_layer.jl (2.3 KB)

Could you kindly point out what I am doing wrong?
Thank you in advance, Stefano

I am using Julia v"1.8.5" with the following packages
ComponentArrays v0.13.7
DifferentialEquations v7.7.0
Lux v0.4.36
Optimization v3.11.2
OptimizationOptimisers v0.1.1
SciMLSensitivity v7.20.0
StableRNGs v1.0.0

2

I tried running your code and didn’t get any error. I used Julia 1.9 and a freshly updated environment with the packages though, so it is probably just something small that has been fixed since your package versions. I would suggest to first just update your packages and see if that is enough, and if not try it with Julia 1.9 and see if that helps.

3

Thank you! I’ll try it