Lux is similar to Flax (Jax) where the parameters are kept in a separate variable from the model definition, and they are passed in on the forward pass. Notably, this design choice allows Lux to accept parameters built with ComponentArrays.jl which can be especially helpful when working with libraries that expect flat vectors of parameters.
Flux lies somewhere between Jax and PyTorch. Like PyTorch, the parameters are stored as part of the model. Unlike traditional PyTorch, Flux has “functional” conventions, e.g. `g = gradient(loss, model)` vs. `loss.backward()`. Similar to Flax, the model is a tree of parameters.