What about Trax (https://github.com/google/trax) and how does it compare with Flax or Haiku?
Interesting that googlers who are supposed to use Tensorflow are now actively developing a new autograd engine and at least three new DL frameworks on top of it. What do you think about this segmentation?