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

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 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.