Tbf JAX is super nice. Even easier than PyTorch in many ways and astonishingly fast. XLA is super powerful.