# Flux: How to create a custom multi-layer model with some parameters shared across layers?

**URL:** <https://discourse.julialang.org/t/flux-how-to-create-a-custom-multi-layer-model-with-some-parameters-shared-across-layers/64190>\
**Category:** Machine Learning\
**Tags:** question\
**Created:** [July 7, 2021, 2:41am UTC](https://discourse.julialang.org/t/flux-how-to-create-a-custom-multi-layer-model-with-some-parameters-shared-across-layers/64190 "2021-07-07T02:41:58Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![IRM](https://avatars.discourse-cdn.com/v4/letter/i/e9c0ed/32.png) [@IRM](https://discourse.julialang.org/u/IRM)\
**Post date:** [July 7, 2021, 2:41am UTC](https://discourse.julialang.org/t/flux-how-to-create-a-custom-multi-layer-model-with-some-parameters-shared-across-layers/64190/1 "2021-07-07T02:41:58Z")

</div>

Hi all,

I’m somewhat new to Julia and Flux, and trying to train a model similar to a standard dense multi-layer neural network, but with sharing of some trainable parameters between layers.

To give a concrete example (not exactly what I want, but its close enough to illustrate the problem I’m facing in Flux):  
The model has as parameters a sequence of matrices A\_l and \Lambda\_l, with the latter diagonal and positive-definite.  
For layers l=1, ... , L-1:  
x\_{l+1} = \sigma (\Lambda\_{l+1}^{-1} A\_l \Lambda\_l x\_l)  
and a final output layer  
y = A\_L \Lambda\_L x\_L

The main difficulty is that each matrix \Lambda\_l for l=2, ..., L appears in both layer l and layer l-1. For this reason I cannot just use Chain, at least as far as I know.

What is the best way of coding this in Flux?

I have tried the code below. It works until the last line which gives the error

> ERROR: Only reference types can be differentiated with `Params`.

I have searched for this error and no solutions I found address exactly this problem. I understand that the problem is related to having a vector of arrays in the struct defining the model. But is there a better way of representing such a structure with a flexible number of layers? Or is there a way of getting Flux to differentiate with respect to the arrays As and ds?

Grateful for any assistance!

```julia

using Flux

mutable struct Multi
    As::Vector{Array{Float64}}
    ds::Vector{Array{Float64}}
end

function (m::Multi)(x)
    L = length(m.As)
    for l = 1:(L-1)
        Λ = diagm(exp.(m.ds[l]))
        V = diagm(exp.(-m.ds[l+1]))
        A = m.As[l]
        x = σ.( V * A * Λ * x )
    end
    Λ = diagm(exp.(m.ds[L]))
    A = m.As[L]
    return A * Λ *x
end
Flux.@functor Multi

m = Multi([randn(nh,ni), randn(no,nh)],[randn(ni), randn(nh)])

x = 0
y = 1

m(1) # check that the model evaluates

function loss(x,y)
    ŷ = m(x)
    sum((y .- ŷ)^2)
end

grads = gradient(() -> loss(x, y), params(m))
grads[1]

```

---

<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:** [July 7, 2021, 4:06am UTC](https://discourse.julialang.org/t/flux-how-to-create-a-custom-multi-layer-model-with-some-parameters-shared-across-layers/64190/2 "2021-07-07T04:06:50Z")

</div>

`Chain` uses a [tuple](https://github.com/FluxML/Flux.jl/blob/master/src/layers/basic.jl#L25), can you do the same?

```julia
struct Multi{T1<:Tuple,T2<:Tuple}
    As::T1
    ds::T2
end

m = Multi((randn(nh,ni), randn(no,nh)), (randn(ni), randn(nh)))

```

AFAICT `Multi` doesn’t need to be mutable either.

---

<div class="post-metadata">

**Author:** ![IRM](https://avatars.discourse-cdn.com/v4/letter/i/e9c0ed/32.png) [@IRM](https://discourse.julialang.org/u/IRM)\
**Post date:** [July 7, 2021, 4:39am UTC](https://discourse.julialang.org/t/flux-how-to-create-a-custom-multi-layer-model-with-some-parameters-shared-across-layers/64190/3 "2021-07-07T04:39:40Z")

</div>

~~Thanks for the suggestion, but unfortunately this gives the same error: Only reference types can be differentiated with `Params` .~~

**Edit:** actually it does seem to work, I was doing something silly before. Thanks for the help!