Custom Flux.jl Layer Not Updating (problem with Flux.trainable?)

I am having a problem getting a custom layer to update, and have tracked the issue to an apparent disconnect between Flux.trainable and Flux.setup. I read through related-sounding posts and issues, which seem to generally be cases of “model won’t learn when gradient is zero”; my issue, however, isn’t just that the layer won’t learn, but rather that Flux explicitly warns that the layer is un-trainable.

Here is my minimal working example:

# trivial layer
mutable struct MWE
(m::MWE)(x) = m.a * x
Flux.@layer MWE

# see that the layer has exactly 1 trainable parameter, `a`
testlayer = MWE(rand())
@assert Flux.trainable(testlayer) == (; a=testlayer.a)

# get error that the layer has NO trainable parameters!
Flux.setup(Flux.Adam(), testlayer)

So, trainable can see the parameter a, but setup somehow misses it. If this layer is used in a larger Chain, the network will learn, but the MWE layer’s parameter will not change.

I can work around this for my current actual use case, but this is concerning and perplexing for me–can anyone point out what’s missing for getting this layer to update?


Classic. The MOMENT I posted the question, I realized the issue.

The parameter a needs to be an array, even if it’s just a single element. That’s literally it.

I’ll leave this post in the hopes that others can learn from my mild embarrassment.

1 Like