JAX is quite popular in many labs outside of Google doing large scale training runs, because up until recently the parallelism ergonomics were way better. PyTorch core is catching up (maybe already witn the latest release, haven’t used it yet) and there are a lot of PyTorch using projects to study though.