Lux.jl icon indicating copy to clipboard operation
Lux.jl copied to clipboard

Allow "const" arrays as inputs to `@compact`

Open avik-pal opened this issue 1 year ago • 0 comments

Currently if we see an array as a kwarg to @compact we put it in trainable mode.

But consider an UDE as a motivating example:

@compact(; model, p_true) do x, p
    dudt(u, p, t) = <nn part> .+ <use p_true>
    ...
end

We could in-principle construct the model and then use freeze, but freeze is known to not work well with non-NamedTuple inputs so it becomes a problem for sciml use cases.

Proposed solution(s):

  • Add a @const or @fixed "pseudo"-macro that doesn't really exist but if we find that symbol during parsing, we make that parameter fixed.
  • Add a reserved kwarg that allows passing a list of symbols that should be considered non-trainable. This is a bit more clunky and not something I prefer to do.

avik-pal avatar Apr 15 '24 22:04 avik-pal