Export GPT2 weights in equinox format
https://docs.kidger.site/equinox/api/serialisation/
explains how it's done.