I'm obviously interested in Jax (otherwise why would I bother commenting on this post?), but it's clearly not mature enough yet for regular DL practitioners. Fo

I'm obviously interested in Jax (otherwise why would I bother commenting on this post?), but it's clearly not mature enough yet for regular DL practitioners. For cutting edge research, especially done at Google, sure. This is fine, it will get better, and I can see myself switching to it in the future. But today, looking at Jax based NN libraries, it's a mess. Stax, Trax, Flax, Haiku, now Objax. It's like the TF mess all over again: old style core TF, Keras, TensorPack, eager mode, estimators, ugh... From what I can see, Flax is the leader (8 months old, 800 commits, last commit is yesterday, similar enough to Pytorch to make the switch easy). So if I wanted to switch to Jax today from Pytorch, why would I choose newborn Objax over Flax? I looked at the documentation and I have no idea how it is more "object oriented" than Flax, or why would I care. The main dev being so bad at marketing his creation leads me to believe it's DOA.