# (Flux/Lux) Custom Layers as Functions of Other Layers

**URL:** <https://discourse.julialang.org/t/flux-lux-custom-layers-as-functions-of-other-layers/97913>\
**Category:** General Usage\
**Tags:** question, machine-learning\
**Created:** [April 25, 2023, 5:50pm UTC](https://discourse.julialang.org/t/flux-lux-custom-layers-as-functions-of-other-layers/97913 "2023-04-25T17:50:28Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![Bizzi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bizzi/32/51484_2.png) [@Bizzi](https://discourse.julialang.org/u/Bizzi)\
**Post date:** [April 25, 2023, 5:50pm UTC](https://discourse.julialang.org/t/flux-lux-custom-layers-as-functions-of-other-layers/97913/1 "2023-04-25T17:50:28Z")

</div>

Hey everyone. For some experiments, I want to set up a network where some of the weights are not trained, but instead are given as functions of other (trained) weights in my network.

For a simple example, take the NN of the following image: I would like to be able to force W3=2\*W1 and W4=3\*W2, then train W1 and W2 normally.

![image](https://global.discourse-cdn.com/julialang/original/3X/7/3/731c3f39b0baeddf6930817ceaf2241784c64029.png)

Is this possible within the SciML environment? I’m quite new to Julia and SciML as a whole, so I honestly wouldn’t even know how to begin. [The Flux page on Custom Layers](https://fluxml.ai/Flux.jl/stable/models/advanced/), opaque as I find it to be, doesn’t seem to consider this possibility.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [April 25, 2023, 6:46pm UTC](https://discourse.julialang.org/t/flux-lux-custom-layers-as-functions-of-other-layers/97913/2 "2023-04-25T18:46:37Z")

</div>

I think in this case using the same layer objects two times each could work?

---

<div class="post-metadata">

**Author:** ![Bizzi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bizzi/32/51484_2.png) [@Bizzi](https://discourse.julialang.org/u/Bizzi)\
**Post date:** [April 25, 2023, 6:56pm UTC](https://discourse.julialang.org/t/flux-lux-custom-layers-as-functions-of-other-layers/97913/3 "2023-04-25T18:56:40Z")

</div>

While that could possibly work, it does not seem to generalize to cases where the mapping is not just the identity, which is what actually interests me. I’ll edit the question to better reflect that!

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [April 25, 2023, 7:29pm UTC](https://discourse.julialang.org/t/flux-lux-custom-layers-as-functions-of-other-layers/97913/4 "2023-04-25T19:29:22Z")

</div>

> [@Bizzi](#):
>
> For a simple example, take the NN of the following image: I would like to be able to force W3=2_W1 and W4=3_W2, then train W1 and W2 normally.

Is your diagram doing something like this?

```julia
struct Diamond{T} # store two matrices
    W1::T
    W2::T
end

Flux.@functor Diamond # make sure Flux can see them

function (d::Diamond)(A) # write out the forward pass
    B = d.W1 * A
    C = d.W2 * A
    D1 = 2 * d.W1 * B
    D2 = 2 * d.W2 * C
    D1 + D2 # assume D is the sum of the two inputs
end

m = Chain(Dense(10=>10, relu), Diamond(randn32(10,10), randn32(10,10)))

m(rand32(10)) # it runs

```

It would be fine to have say `D1 = 2 * (d.W1 .^ 2) * B`, or some other function of `W1` before using it a second time.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [April 25, 2023, 7:50pm UTC](https://discourse.julialang.org/t/flux-lux-custom-layers-as-functions-of-other-layers/97913/5 "2023-04-25T19:50:31Z")

</div>

I’m not sure what it would look like in Flux, but in Lux you could just pass identical or modified versions of the parameter objects to both layers