# Flux params restrictive

**URL:** <https://discourse.julialang.org/t/flux-params-restrictive/77428>\
**Category:** General Usage\
**Tags:** performance, flux, zygote, garbage-collection\
**Created:** [March 4, 2022, 10:06pm UTC](https://discourse.julialang.org/t/flux-params-restrictive/77428 "2022-03-04T22:06:50Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![cortner](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cortner/32/204_2.png) [@cortner](https://discourse.julialang.org/u/cortner)\
**Post date:** [March 4, 2022, 10:06pm UTC](https://discourse.julialang.org/t/flux-params-restrictive/77428/1 "2022-03-04T22:06:50Z")

</div>

I’m working on integrating the `Zygote.Params` idea into my codes to simplify parameter wrangling, and I was struck by the fact that

```nohighlight
params!(p::Params, x::AbstractArray{<:Number}, seen = IdSet()) = push!(p, x)

```

only arrays of numbers are allowed as parameter fields, but not arrays of arrays, in particular I’m interested in using e.g. `Vector{SVector{N, T}}`. I get it is equivalent to matrix, but this is just an example. In fact, I’m even more interested in `Array{N, T}` where `T` is an abstract array such as `struct MyT a::Float64; b::SVector{3, Float64} end`

Is there a strong reason to only allow `x::AbstractArray{<:Number}`? What are the pitfalls if I try to extend this behaviour?

---

<div class="post-metadata">

**Author:** ![cortner](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cortner/32/204_2.png) [@cortner](https://discourse.julialang.org/u/cortner)\
**Post date:** [March 5, 2022, 12:06am UTC](https://discourse.julialang.org/t/flux-params-restrictive/77428/2 "2022-03-05T00:06:35Z")

</div>

hm … I was just told in a Flux issue that Flux is in fact trying to drop implicit parameters altogether. Is there a discussion somewhere of why and what will it be replaced with?

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [March 5, 2022, 1:20am UTC](https://discourse.julialang.org/t/flux-params-restrictive/77428/3 "2022-03-05T01:20:09Z")

</div>

I’m not sure of the original intent, but one thing this allows is the use of Vectors for holding a list of layers. If it treated all arrays as parameters, the [`AbstractVector` version of `Chain`](https://github.com/FluxML/Flux.jl/blob/f49e81e7b366689bc57a65f8817de4353f144270/src/layers/basic.jl#L35) wouldn’t work, for example.

As an easy workaround, you can add params yourself to the collection by `push!`ing them manually.

---

<div class="post-metadata">

**Author:** ![cortner](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cortner/32/204_2.png) [@cortner](https://discourse.julialang.org/u/cortner)\
**Post date:** [March 5, 2022, 1:25am UTC](https://discourse.julialang.org/t/flux-params-restrictive/77428/4 "2022-03-05T01:25:34Z")

</div>

Ok but I don’t want to start using a paradigm that’s in the way out anyhow. What Should I do instead? How should I structure my models?

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [March 5, 2022, 1:28am UTC](https://discourse.julialang.org/t/flux-params-restrictive/77428/5 "2022-03-05T01:28:48Z")

</div>

“On the way out” is not going to happen for some months yet, and I imagine in the meantime you’ll want to get work done 😉 . We’ll make sure to provide a migration path when the time comes.

---

<div class="post-metadata">

**Author:** ![cortner](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cortner/32/204_2.png) [@cortner](https://discourse.julialang.org/u/cortner)\
**Post date:** [March 5, 2022, 3:14am UTC](https://discourse.julialang.org/t/flux-params-restrictive/77428/6 "2022-03-05T03:14:50Z")

</div>

thanks for answering me in two places at once. I want to run one more test and will then maybe have one more clarifying question.

---

<div class="post-metadata">

**Author:** ![cortner](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cortner/32/204_2.png) [@cortner](https://discourse.julialang.org/u/cortner)\
**Post date:** [March 8, 2022, 6:07am UTC](https://discourse.julialang.org/t/flux-params-restrictive/77428/7 "2022-03-08T06:07:28Z")

</div>

@ToucheSir – thanks for your comments so far; I now have one more question and would be very grateful for your input. An important piece of background information here is that I am not planning to _use_ standard ML models through flux, but I’m working a fairly specialized class of models related to equivariant message passing, and am really writing all components from scratch. However, I can see that in the future I might leverage some of GeomtricFlux or even use a few Flux components to give additional flexibility to my models.

My alternative plan to using the implicit parameters of Flux / Zygote is to specify parameters through a nested `NamedTuple`, e.g., like this:

```nohighlight
mutable struct ACE{T, TB}
   basis::TB 
   c::Vector{T}
end

params(ace::ACE) = (basis = params(ace.basis), c = ace.c)

function set_params!(ace::ACE, p::NamedTuple)
   set_params!(ace.basis, p.basis)
   ace.c = p.c
   return ace 
end 

```

In this case, a Zygote call such as

```nohighlight
p0 = params(ace)
g = Zygote.gradient( p -> somefunction( set_params!(ace, p), X ), p0 )

```

will return a gradient `g::NamedTuple` with precisely the same structure as the `p0` params NamedTuple. I then wrap them into a type that treats them as abstract vectors so I can add, multiply etc and do optimization.

I’ve implemented a prototype for this and it works quite well. Going a step further, I’ve also experimented with specifying rrules where the gradient w.r.t. the model is the gradient w.r.t. parameters. This has worked equally well. All this takes a little more overhead than the implicit parameters, but since I have to write custom rrules for all my model components anyhow I’m still considering it.

Arguably one advantage of this approach is that I’m not limited to using only Arrays for storing parameters but have much more flexibility.

Do you have any thoughts on this approach and how well it might interact with the future plans for Flux, or is your advise still to stick with implicit parameters for now and then transition once Flux transitions?

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [March 8, 2022, 5:06pm UTC](https://discourse.julialang.org/t/flux-params-restrictive/77428/8 "2022-03-08T17:06:33Z")

</div>

I would recommend against overriding `params` directly since that causes confusion with Flux’s definition. Case in point, you’re actually using “explicit” params here since `p0` is a NamedTuple and not a `Params`. So assuming you’re getting the gradients you want, this should be reasonably future-proof.

---

<div class="post-metadata">

**Author:** ![cortner](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cortner/32/204_2.png) [@cortner](https://discourse.julialang.org/u/cortner)\
**Post date:** [March 8, 2022, 8:27pm UTC](https://discourse.julialang.org/t/flux-params-restrictive/77428/9 "2022-03-08T20:27:33Z")

</div>

Thank you!

Just to make sure I understand correctly - you don’t see a problem with the NamedTuple style explicit parameters and this might be more future proof than going back to implicit Flux paramters?

But you are suggesting to name `params` it something else so it doesn’t interfere with Flux.

---

<div class="post-metadata">

**Author:** ![cortner](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cortner/32/204_2.png) [@cortner](https://discourse.julialang.org/u/cortner)\
**Post date:** [March 9, 2022, 2:45pm UTC](https://discourse.julialang.org/t/flux-params-restrictive/77428/10 "2022-03-09T14:45:01Z")

</div>

Thank you again - you’ve been most helpful!
