No really, it force you into pure, functional patterns. Frameworks build on top of JAX have done a lot of experiments to find a good way to deal with states while keeping the code clean and not too magicy.

Their is no clear best pattern at the moment. My favorite would be Flax and their recent Linen API which is a refined effort that pays off when using their framework.

No one does. (Also, that's a flax story, not a jax story. A bit confusing, but think of jax like numpy as flax like pytorch -- jax is the low level, flax is the high level.)

I haven't found a way that I like. The closest I've come is tftorch: https://twitter.com/theshawwn/status/1311925180126511104

I think it's important to have global scope names. "biggan.discriminator.3.conv1.kernel.b" is perfectly sensible: it's the bias value of the kernel for the convolution of the third block of your biggan discriminator.

Everyone tries to treat model variables as interchangeable nameless parts. I hate it. Every variable has a global name, conceptually.

A global name also has other benefits. It becomes far easier to create EMA weights, for example, since you can filter the variables by name.