Very good that Pytorch emerged as a serious contender to TF. While TF still provided more production grade tools (TFX, TensorRT, TF serving), Pytorch continue to evolve and hope soon we have a more complete ecosystem
It seems that the JAX developers are focusing their time on making the core framework better and are leaving the task of building high-level APIs to the community for now. I suspect we'll see a few high-level APIs emerge over the next few months that explore different approaches before the community settles on a particular one.