What is more some of the deep learning libraries on top of JAX (like flax, linen) are close to a one-to-one copy of PyTorch. However the underlying implementation is in some ways more interesting than the one of PyTorch and naturally accomodates higher derivatives and also the way that optimization and state initialization is implemented is more functional and "theoretically sound".