This is a simple benchmark of a transformer language model implemented in JAX and PyTorch. See blog post: http://www.echonolan.net/posts/2021-09-06-JAX-vs-PyTorch-A-Transformer-Benchmark.html