Implement LoRA in JAX/Equinox

Implement LoRA in such a way that it can be used with our GPT2 implementation.