Building a Tiny Transformer for Addition in JAX