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

Package Announcements
, ,
21

I did a major overhaul of the Lux documentation https://lux.csail.mit.edu/. I would appreciate feedback regarding the new look and the unified approach. Some improvements:

  • Lists in the previous docs were broken due to how the generated markdown was printed (mkdocs seems to have a slightly different way of handling new lines). That is fixed now.
  • There were complaints regarding how each docstring was not separated and hence hard to read. Now the docstrings have proper bounding boxes
    image
    image711×222 9.4 KB
  • Finally all the docs are in the same place which means they are easy to search!
  • The versioning is at least visible now. I couldn’t work it out for mkdocs material previously. Though the versioning is not great (vitepress doesn’t natively support versioning)

Known Issues:

  • Only dev is live. This is a problem from TagBot which I am looking into.
  • Some docstrings have !!! warning instead of nice warnings/info, etc blocks. I need to go through all the packages and update their docstrings
22

Wow, it looks great!

There is an import error in the first example of this section: https://lux.csail.mit.edu/dev/tutorials/beginner/Basics/main#Automatic%20Differentiation

Screenshot 2023-08-21 at 20.33.16
Screenshot 2023-08-21 at 20.33.16734×452 50.3 KB

23

These look amazing! The homepage is particularly impressive.

One minor complaint I have is that the “Getting Started” page is quite sparse - aside from just getting Julia + Lux installed, we should also show some examples of how to use Lux in practice; if you need to be shown how to install Julia or a package, you probably also need to be shown how to really use that package :slightly_smiling_face:

We could show a simple example of creating and training a really basic NN, and then provide links to more advanced or domain-specific examples for a variety of use cases.

24

Great work! Always inspired by your amazing work. These docs should be the new standard :D.

As for the looks, so far I only have one suggestion:

it would be good to highlight in which section we are currently at, on the right panel. Current behaviour doesn’t show anything.

25

image|50%
image|50%429×503 20.3 KB

On the tutorials page of this whole box only the </> is clickable, it’s hard to even notice that bit.

26

Can you elaborate on the tools you used to build the Lux website and docs?

27

Thanks for all the suggestions. I have aggregated them in an issue Documentation Enhancement Suggestions · Issue #387 · LuxDL/Lux.jl · GitHub.

Can you elaborate on the tools you used to build the Lux website and docs?

It uses

28

why vite? how the future looks like for this option, it’s only 3 years old. It’s definitely nicer at first glance, but beyond that? Just curious about the motivation for such change, if there is any :smiley: .

29

This might be overkill, but a brief mention of the f0 and how it relates to Float32 might be helpful in this line

Amazing docs though, this template is fantastic!

30

(My entire JS knowledge comes from 2 lab assignments I had to do in my undergrad, so take everything I say with a grain of salt)

I was already kind of familiar with Vuejs, and it seems vuepress is in maintenance mode (GitHub - vuejs/vuepress: 📝 Minimalistic Vue-powered static site generator), and the current effort has shifted towards vitepress (with it being hosted in the official vuejs organization).

The other option was docusaurus, but I did not have any examples in the julia ecosystem to go off of here. There’s https://expronicon.rogerluo.dev/ which has been using vitepress for quite some time now.

31

Is there anyway in Lux to use Float64 for the networks parameters like in Flux? I couldnt quite figure this out from the documentation.

32

Just change the weights to Float64.

33

I also found that surprisingly inconvenient to do actually. (Unless it’s changed recently)

34

I see; I can add a function like f64/f32/f16 to automatically perform this conversion.

35

Or just a kwarg when constructing the layer

36

Having a kwarg for every layer becomes infeasible very quickly. FWIW, in the latest release, constructs the parameters with the correct types if you do init_weight=glorot_uniform(Float64) and similar for different layers.

37

oh that’s all I need, thank you. I think that didn’t exist a few versions ago, correct?

38

Yeah pretty much. It got merged yesterday Allow specifying the return eltype of the arrays by avik-pal · Pull Request #4 · LuxDL/WeightInitializers.jl · GitHub

39

Having a kwarg for every layer becomes infeasible very quickly.

I’ve been experimenting with a “generative context” to pass defaults for device, numeric precision and RNG through a computation:

Maybe something like that?

40

Using ComponentArrays this should do the trick

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