# \[ANN\] Lux.jl: Explicitly Parameterized Neural Networks in Julia

**URL:** <https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689>\
**Category:** Package Announcements\
**Tags:** package, announcement, machine-learning\
**Created:** [May 25, 2022, 11:24pm UTC](https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689 "2022-05-25T23:24:01Z")\
**Posts on this page:** 11\
**Page:** 3

<div class="post-metadata">

**Author:** ![TheLateKronos](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/thelatekronos/32/12824_2.png) [@TheLateKronos](https://discourse.julialang.org/u/TheLateKronos)\
**Post date:** [August 26, 2023, 1:27pm UTC](https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689/41 "2023-08-26T13:27:30Z")

</div>

> [@bertschi](#):
>
> ```julia
> ps, st = Lux.setup(rng, model)
> ps = Float64.(ComponentVector(ps))
> 
> ```

Wound it not be sufficient to do

```julia
ps, st = Lux.setup(rng, model)
ps = Float64.(ps)

```

?

---

<div class="post-metadata">

**Author:** ![bertschi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bertschi/32/33462_2.png) [@bertschi](https://discourse.julialang.org/u/bertschi)\
**Post date:** [August 26, 2023, 2:03pm UTC](https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689/42 "2023-08-26T14:03:50Z")

</div>

That would work if `ps` were a flat vector, but `Lux` uses nested named tuples for its parameters.

Seems that the state `st` is even more tricky as it contains leafs of different types, e.g., `training = Val{true}` besides `Vector{Float32}` parameters. Thus, `ComponentVector(st)` will not work. On the other hand, as soon as you run the model a single time via `y, st = Lux.apply(model, x, ps, st)` the state will be promoted to the type of `x` and `ps` anyways, so you can use that to convert instead.

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [August 28, 2023, 2:42pm UTC](https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689/43 "2023-08-28T14:42:46Z")

</div>

The latest stable release has the functions available [Utilities | LuxDL Docs](https://lux.csail.mit.edu/dev/api/Lux/utilities#updating-floating-point-precision)

---

<div class="post-metadata">

**Author:** ![Ashwani\_Rathee](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ashwani_rathee/32/49551_2.png) [@Ashwani\_Rathee](https://discourse.julialang.org/u/Ashwani_Rathee)\
**Post date:** [September 6, 2023, 5:15pm UTC](https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689/44 "2023-09-06T17:15:15Z")

</div>

I love the style and theme of the website! [Getting Started | LuxDL Docs](https://lux.csail.mit.edu/dev/introduction/) could be more descriptive though, given its an introduction despite focusing on the coding aspect which is alright, I think more useful and descriptive comments could be added.

---

<div class="post-metadata">

**Author:** ![braamvandyk](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/braamvandyk/32/5086_2.png) [@braamvandyk](https://discourse.julialang.org/u/braamvandyk)\
**Post date:** [September 7, 2023, 6:19am UTC](https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689/45 "2023-09-07T06:19:14Z")

</div>

I started looking at Lux based on this thread, and I must say that I am truly appreciative of the documentation. Thanks to everyone who put in the time to work on the package and the website!

In the spirit of giving constructive feedback, there is just one proposal I would like to make. The font used for code blocks makes minimal distinction between a period and a comma. One example from the very first tutorial in “Julia and Flux for the Uninitiated” in the section on “(Im)mutability”:

![image](https://global.discourse-cdn.com/julialang/original/3X/b/d/bd5d1be96ff42fc4c28da3bce05e4083f1916421.png)

If you know any Julia, this won’t be a problem, but this tutorial is specifically aimed at people who don’t, which may result in a little (or a lot of) head-scratching.

For comparison, when I copy and paste the code to my REPL:

![image](https://global.discourse-cdn.com/julialang/original/3X/0/f/0fded40cb29a419eaa8d83b63cc1eba3b770e85e.png)

In any event, I shall be directing my team members interested in learning Julia and deep-learning to the Lux docs. It does an infinitely better job at explaining things than I have been doing.

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [April 8, 2024, 3:47pm UTC](https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689/46 "2024-04-08T15:47:31Z")

</div>

## Some New Major Updates (till v0.5.33)

1. Lux now has in-built distributed training support via MPI – [Distributed Data Parallel Training | Lux.jl Docs](https://lux.csail.mit.edu/stable/manual/distributed_utils). It is effectively a rewrite of my older package [GitHub - avik-pal/FluxMPI.jl: Distributed Data Parallel Training of Deep Neural Networks](https://github.com/avik-pal/FluxMPI.jl) (that is now archived), but allows NVIDIA GPU communication via NCCL!
2. SimpleChains is available as a backend ([https://lux.csail.mit.edu/stable/api/Lux/switching\_frameworks#Lux.ToSimpleChainsAdaptor](https://lux.csail.mit.edu/stable/api/Lux/switching_frameworks#Lux.ToSimpleChainsAdaptor)) for small neural networks, which means you can write everything in Lux and still use SimpleChains with just 1 additional line of code ([MNIST Classification with SimpleChains | Lux.jl Docs](https://lux.csail.mit.edu/stable/tutorials/beginner/4_SimpleChains))

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/224680_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [April 9, 2024, 6:03am UTC](https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689/47 "2024-04-09T06:03:39Z")

</div>

That’s awesome! IIRC, one of SimpleChains.jl’s main selling points is that it can be allocation-free: how does this pair with Lux’s purely functional style?

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [April 9, 2024, 3:18pm UTC](https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689/48 "2024-04-09T15:18:24Z")

</div>

It still maintains the purity in the sense of no-side effects and same inputs =\> same outputs. But yeah you can’t get full performance of simple chains staying in the lux api or for that matter using the chainrules api. FWIW most simple chains users from sciml point would use the chain rules api (See [Faster Neural Ordinary Differential Equations with SimpleChains · SciMLSensitivity.jl](https://docs.sciml.ai/SciMLSensitivity/stable/examples/neural_ode/simplechains/)) and not the fully non-allocating `train_batched!` API.

Lux API is simpler to use and more people are familiar with it (since we mimicked the flux api for layers). SimpleChains is extremely fast but the API is somewhat non-traditional. So now users can write the model in Lux and (if possible) convert it to simple chains (getting the performance boost) and never have to learn the details of how to write in SimpleChains

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/224680_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [April 9, 2024, 3:57pm UTC](https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689/49 "2024-04-09T15:57:52Z")

</div>

Perhaps SimpleChains would be a good candidate for a DifferentiationInterface binding, which could then be accessed from Lux

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [April 16, 2024, 4:03pm UTC](https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689/50 "2024-04-16T16:03:50Z")

</div>

## New Additions (v0.5.36)

- [`DynamicExpressionsLayer`](https://lux.csail.mit.edu/stable/api/Lux/interop#Lux.DynamicExpressionsLayer) lets you now embed Symbolic Expressions into Lux models. So we can fit a model using SymbolicRegression.jl and swap it into a lux chain (see this new tutorial [https://lux.csail.mit.edu/stable/tutorials/advanced/2\_SymbolicOptimalControl](https://lux.csail.mit.edu/stable/tutorials/advanced/2_SymbolicOptimalControl)). (Thanks `@milescranmer` for a starter code on this)

- `@compact` is now the recommended API for creating custom models. It part of the stable API (as of `v0.5.35`), and its capabilities have been extended to allow accessing parameters inside the model. For example, writing a UDE is now as simple as:

```julia
ude = @compact(; nn, p_true, solver=Tsit5(), tspan=(0.0,1.0),
        kwargs...) do x, ps
    # Just an arbitrary UDE
    dudt(u, p, t) = nn(x, p.nn) .+ sum(p.p_true)
    prob = ODEProblem{false}(dudt, x, tspan, ps)
    return solve(prob, solver; kwargs...)
end

```

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [April 27, 2024, 5:33pm UTC](https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689/51 "2024-04-27T17:33:45Z")

</div>

Split this topic on April 27, 2024, 5:33pm UTC

20 posts were split to a new topic: [Nested AD with Lux etc](https://discourse.julialang.org/t/nested-ad-with-lux-etc/113573)

[Previous page](https://discourse.julialang.org/t/ann-lux-jl-explicitly-parameterized-neural-networks-in-julia/81689.md?page=2)