Export GPT2 weights in equinox format

https://docs.kidger.site/equinox/api/serialisation/ explains how it's done.