Haiku/Jax implementation of transformer example from https://github.com/deepmind/dm-haiku/tree/main/examples/transformer