JAX vs Tensorflow vs Pytorch: Building a Variational Autoencoder (VAE)

A side-by-side comparison of JAX, Tensorflow and Pytorch while developing and training a Variational Autoencoder from scratch