flax_transformers Some basic transformer architectures in flax. XLA compiler go brrrrr Leftovers from my tinkering. May need some debugging still.