# Building a custom Layer in Lux.jl

**URL:** <https://discourse.julialang.org/t/building-a-custom-layer-in-lux-jl/92091>\
**Category:** General Usage\
**Created:** [December 24, 2022, 4:06pm UTC](https://discourse.julialang.org/t/building-a-custom-layer-in-lux-jl/92091 "2022-12-24T16:06:26Z")\
**Posts on this page:** 11\
**Page:** 1

<div class="post-metadata">

**Author:** ![erlebach](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/erlebach/32/12973_2.png) [@erlebach](https://discourse.julialang.org/u/erlebach)\
**Post date:** [December 24, 2022, 4:06pm UTC](https://discourse.julialang.org/t/building-a-custom-layer-in-lux-jl/92091/1 "2022-12-24T16:06:26Z")

</div>

I am constructing a custom layer, and have an error:

```julia
ERROR: MethodError: no method matching iterate(::Type{Vector{Float32}})
Closest candidates are:
  iterate(::Union{LinRange, StepRangeLen}) at range.jl:872
  iterate(::Union{LinRange, StepRangeLen}, ::Integer) at range.jl:872
  iterate(::T) where T<:Union{Base.KeySet{<:Any, <:Dict}, Base.ValueIterator{<:Dict}} at dict.jl:712
  ...
Stacktrace:
 [1] Polylayer(coefs::Vector{Float32})
   @ Main ~/src/2022/basic_UODE/custom_lux_layer/poly_layer.jl:13
 [2] top-level scope
   @ ~/src/2022/basic_UODE/custom_lux_layer/poly_layer.jl:17

```

which should be easy to decipher, but I can’t figure it out. I created a MWE, which is:

```julia
using Lux

struct Polylayer{C} <: Lux.AbstractExplicitLayer
    coefs::C
end

function Polylayer(coefs)
    dtype = (typeof(coefs))
    return Polylayer{dtype...}(coefs)
end

a = Vector{Float32}(undef, 4)
pol = Polylayer(a)

```

Any insight it appreciated.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 24, 2022, 4:32pm UTC](https://discourse.julialang.org/t/building-a-custom-layer-in-lux-jl/92091/2 "2022-12-24T16:32:37Z")

</div>

For reference the docs are:

[http://lux.csail.mit.edu/stable/manual/interface/#singular-layer](http://lux.csail.mit.edu/stable/manual/interface/#singular-layer)

and you’re missing a few functions from that.

But your immediate issue is that the constructor function doesn’t quite make sense? What are you trying to do? You can just delete

```julia
function Polylayer(coefs)
    dtype = (typeof(coefs))
    return Polylayer{dtype...}(coefs)
end

```

and it’ll be fine.

---

<div class="post-metadata">

**Author:** ![erlebach](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/erlebach/32/12973_2.png) [@erlebach](https://discourse.julialang.org/u/erlebach)\
**Post date:** [December 24, 2022, 4:43pm UTC](https://discourse.julialang.org/t/building-a-custom-layer-in-lux-jl/92091/3 "2022-12-24T16:43:28Z")

</div>

Hi,

Yes, I know there are functions missing. But since I am only testing the layer by calling the functions manually, I did not think the added functions were necessary. I also created a `show` function, thinking that would do the trick. Finally, I removed other functionality to create a MWE.

I want to construct a custom layer that implements a polynomial of specified degree for a pair of variables. So I created the structure:

```julia
struct Polylayer{C} <: Lux.AbstractExplicitLayer
    in_dims::Int
    out_dims::Int
    degree::Int
    coefs::C
end

```

and the following constructors:

```julia
function Polylayer(in_dims::Int, out_dims::Int, coefs)
    degree = length(coefs)
    dtype = (typeof(coefs))
    return Polylayer{dtype...}(in_dims, out_dims, degree, coefs)
end

function Polylayer(in_dims::Int, out_dims::Int, degree::Int)
    coefs = zeros(degree) 
    dtype = (typeof(coefs))
    #println(coefs, typeof(coefs))
    return Polylayer{dtype...}(in_dims, out_dims, degree, coefs)
end

```

I am modeling this approach on the `Dense` layer found in the `Lux.jl` source code. I looked for examples of custom layers for `Lux.jl` online and there are very few of them at this time.

I know I must implement `initialparameters`, `parameterlength` and `statelength`.

The layer will receive one scalar variable, and output the evaluation of the polynomial `Polylayer(x, p, st)`. My objective is to have Lux train the polynomial coefficients. The larger experiment is an execution of the LV equations where I express the nonlinear term directly as a polynomial and use UODE to find the coefficients. This is a test example for a future larger problem.

---

<div class="post-metadata">

**Author:** ![erlebach](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/erlebach/32/12973_2.png) [@erlebach](https://discourse.julialang.org/u/erlebach)\
**Post date:** [December 24, 2022, 5:08pm UTC](https://discourse.julialang.org/t/building-a-custom-layer-in-lux-jl/92091/4 "2022-12-24T17:08:24Z")

</div>

I implemented a more complete layer, and it does generate errors.

```julia
struct Polylayer{C,F} <: Lux.AbstractExplicitLayer
    in_dims::Int
    out_dims::Int
    degree::Int
    coefs::C
    init_weight::F
end

function Base.show(io::IO, d::Polylayer)
    println(io, "Polylayer($(d.in_dims) => $(d.out_dims), degree: $(d.degree)")
    println("inside show")
end

function Polylayer(in_dims::Int, out_dims::Int, coefs, init_weight=Lux.rand32)
    degree = length(coefs)
    dtype = (typeof(coefs), typeof(init_weight))
    return Polylayer{dtype...}(in_dims, out_dims, degree, coefs, init_weight)
end

function initialparameters(rng::AbstractRNG, d::Polylayer)
    return (d.init_weight(rng, d.degree+1))
end

statelength(d::Polylayer) = 0
function parameterlength(d::Polylayer)
    return d.degree + 1
end

```

I just did not think that the additional two functions were necessary, even for minimal testing.  
Rather, I thought the additional functions would be necessary when actually using the layer within a chain.

Thanks!

---

<div class="post-metadata">

**Author:** ![erlebach](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/erlebach/32/12973_2.png) [@erlebach](https://discourse.julialang.org/u/erlebach)\
**Post date:** [December 24, 2022, 5:23pm UTC](https://discourse.julialang.org/t/building-a-custom-layer-in-lux-jl/92091/5 "2022-12-24T17:23:28Z")

</div>

OK, I now understand. No need to include the weight arrays in the struct. Rather when I call `Lux.setup`, the parameters are initialized.

I must say that I still have difficulty wrapping my head around the Julia way of coding. Yes, Julia is extremely efficient, but the approach to coding takes lots of getting used to. I am not convinced it is faster than conventional languages. So far, the development process (for me) is slower. But of course, that is probably due to the many years of coding in Fortran, C++, Python.

Cheers, and happy holidays!

---

<div class="post-metadata">

**Author:** ![erlebach](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/erlebach/32/12973_2.png) [@erlebach](https://discourse.julialang.org/u/erlebach)\
**Post date:** [December 24, 2022, 5:46pm UTC](https://discourse.julialang.org/t/building-a-custom-layer-in-lux-jl/92091/6 "2022-12-24T17:46:47Z")

</div>

@ChrisRackauckas,

Here is my next to final custom layer. All seems to work except the very last `Lux.setup(..)` call. Note that if I call `initialparameters()` directly, the function returns a parameterized array. Strange. Here is the code:

```julia
using Lux
using Random

"""
Apply a polynomial to an input Float
The coefficients are trainable. 
First attempt: polynomial of a specified `degree`
"""

struct Polylayer{F} <: Lux.AbstractExplicitLayer
    in_dims::Int
    out_dims::Int
    degree::Int
    init_weight::F
end

function Base.show(io::IO, d::Polylayer)
    println(io, "Polylayer($(d.in_dims) => $(d.out_dims), degree: $(d.degree)")
    println("inside show")
end

function Polylayer(in_dims::Int, out_dims::Int, degree::Int, init_weight=Lux.rand32)
    dtypes = (typeof(init_weight))
    return Polylayer{dtypes}(in_dims, out_dims, degree, init_weight)
end

function initialparameters(rng::AbstractRNG, d::Polylayer)
    return (d.init_weight(rng, d.degree+1))
end

statelength(d::Polylayer) = 0

function parameterlength(d::Polylayer)
    return d.degree + 1
end
# ====================================

rng = Random.default_rng()
model = l.Polylayer(5, 3, 3)
ps, st = Lux.setup(rng, model)

ps # I get an empty parameter list

```

---

<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:** [December 24, 2022, 6:04pm UTC](https://discourse.julialang.org/t/building-a-custom-layer-in-lux-jl/92091/7 "2022-12-24T18:04:25Z")

</div>

> [@erlebach](#):
>
> OK, I now understand. No need to include the weight arrays in the struct. Rather when I call `Lux.setup`, the parameters are initialized.
> 
> I must say that I still have difficulty wrapping my head around the Julia way of coding. Yes, Julia is extremely efficient, but the approach to coding takes lots of getting used to.

I think here the disconnect may be less the language and more the specific pattern the library uses. It helps if you’ve seen [GitHub - google/flax: Flax is a neural network library for JAX that is designed for flexibility.](https://github.com/google/flax/) or [GitHub - google-deepmind/dm-haiku: JAX-based neural network library](https://github.com/deepmind/dm-haiku) before, both of which separate layers from their actual parameters as Lux does. Ordinarily in Julia, one would include such arrays in the layer structs themselves (this is what Flux, Knet and co. do, for example) as you were expecting to be the case.

---

<div class="post-metadata">

**Author:** ![erlebach](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/erlebach/32/12973_2.png) [@erlebach](https://discourse.julialang.org/u/erlebach)\
**Post date:** [December 24, 2022, 6:09pm UTC](https://discourse.julialang.org/t/building-a-custom-layer-in-lux-jl/92091/8 "2022-12-24T18:09:40Z")

</div>

Very true, @ToucheSir ! Intellectually I know how it works, but out of habit, I return to the standard way of doing things. I tried learning Jax recently but have not put in the time. I am a Pytorch person myself.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 24, 2022, 6:14pm UTC](https://discourse.julialang.org/t/building-a-custom-layer-in-lux-jl/92091/9 "2022-12-24T18:14:26Z")

</div>

```julia
function initialparameters(rng::AbstractRNG, d::Polylayer)
    return (coeffs = d.init_weight(rng, d.degree+1),)
end

Lux.initialstates(::AbstractRNG, ::Polylayer) = NamedTuple()

function (l::Polylayer)(x::AbstractMatrix, ps, st::NamedTuple)
    y = ps.coeffs *. x
   ... # do whatever your layer is
end

```

---

<div class="post-metadata">

**Author:** ![erlebach](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/erlebach/32/12973_2.png) [@erlebach](https://discourse.julialang.org/u/erlebach)\
**Post date:** [December 24, 2022, 6:36pm UTC](https://discourse.julialang.org/t/building-a-custom-layer-in-lux-jl/92091/10 "2022-12-24T18:36:32Z")

</div>

Thanks, but I still get empty arrays from `Lux.setup`.  
I noticed you added a comma in `initial parameters`, but it had no effect. I have no idea there is a tiny mistake somewhere. Julia is rather unforgivable, as is any compiled language.

So make sure we are looking at the same code, I reattach it:

```julia
sing Lux
using Random

"""
Apply a polynomial to an input Float
The coefficients are trainable. 
First attempt: polynomial of a specified `degree`
"""

struct Polylayer{F} <: Lux.AbstractExplicitLayer
    out_dims::Int
    degree::Int
    init_weight::F
end

function Base.show(io::IO, d::Polylayer)
    println(io, "Polylayer(out_dims: $(d.out_dims)), degree: $(d.degree)")
end

function Polylayer(out_dims::Int, degree::Int, init_weight=Lux.rand32)
    dtypes = (typeof(init_weight))
    return Polylayer{dtypes}(out_dims, degree, init_weight)
end

function (l::Polylayer)(x::AbstractMatrix, ps, st::NamedTuple)
    y = x
    return y, st
end

function initialparameters(rng::AbstractRNG, d::Polylayer)
     return (coeffs=d.init_weight(rng, d.degree+1),)
end

Lux.initialstates(::AbstractRNG, ::Polylayer) = NamedTuple()

statelength(d::Polylayer) = 0

function parameterlength(d::Polylayer)
    return d.degree + 1
end

# ====================================

rng = Random.default_rng()
Random.seed!(rng, 0)

model = Polylayer(5, 3, Lux.rand32)
ps, st = Lux.setup(rng, model) # <<<<<< produces two empty NamedTuples.
initialparameters(rng, model) # Initializes correctly. 

```

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 24, 2022, 11:03pm UTC](https://discourse.julialang.org/t/building-a-custom-layer-in-lux-jl/92091/11 "2022-12-24T23:03:25Z")

</div>

> [@erlebach](#):
>
> I noticed you added a comma in `initial parameters`, but it had no effect.

That’s on purpose: it should return a tuple instead of an array. See `(1)` vs `(1,)`.

And the function is shadowing instead of extending. You want in that script:

```julia
function Lux.initialparameters(rng::AbstractRNG, d::Polylayer)
     return (coeffs=d.init_weight(rng, d.degree+1),)
end

Lux.initialstates(::AbstractRNG, ::Polylayer) = NamedTuple()

Lux.statelength(d::Polylayer) = 0

function Lux.parameterlength(d::Polylayer)
    return d.degree + 1
end

```