My path to learn jax.
Colab link. This is using flax. (The pytorch counterpart is here)
I will try to take a look at Haiku. Haiku seems very similar to flax.linen. I will take a look at trax. It is OO, and it is more actively developed than objax.
Flax is pretty similar to tensorflow 1.0 I feel like.
I hate google just have all these different high-level packages from different places. flax, trax, haiku etc...... (Saw a discussion here).