3 ms·
From people with more experience than me: does JAX have a good story about how to manage state cleanly, where it needs to occur?
by taliesinb 6y ago
From people with more experience than me: does JAX have a good story about how to manage state cleanly, where it needs to occur?
- sillysaurusx 6y agoNo 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 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.
- nestorD 6y agoNo 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.