JAX is focusing more on performance and probably will be next big think for google machine learning teams like https://github.com/google/flax (most of the people hate tensorflow) Dunno about zygote future