Putting my JAX-trained models on the Hugging Face Hub

Writing the post that I wished I'd found when I started learning whatever it was...

I hadn't uploaded the models that I trained using JAX to the Hugging Face Hub because Transformers has been PyTorch-only since version 5 (though they say they're working to add interoperability with JAX in the future), so it would have been tough to get them working natively with AutoModelForCausalLM and the like.

But then it dawned on me that I'd already written a conversion script that could take my JAX safetensors files and convert them into ones compatible with my PyTorch code. It's actually those converted models that I use for my evals -- so I could use my existing PyTorch script to upload them.

So, I've now uploaded PyTorch-compatible versions of all of my JAX-trained models:

"Writing an LLM from scratch, part 34b -- from bigrams to GPT-2, one component at a time (in JAX)"

• gpjt/jax-no-mha-bias-no-dropout -- the first full LLM trained in the post, in the "Adding LayerNorm" section.
• gpjt/jax-no-mha-bias-with-dropout -- the second full LLM trained in the post, in the "Dropout" section.
• gpjt/jax-with-mha-bias-no-dropout -- the third full LLM trained in the post, in the "Adding bias to the MHA output projections" section.

"Why do OpenAI's GPT-2 weights beat mine? Part three: testing overtraining"

• gpjt/jax-with-mha-bias-no-dropout-extended -- the single-epoch, double-Chinchilla-tokens model.
• gpjt/jax-with-mha-bias-no-dropout-2-epoch -- the model trained on two epochs over the Chinchilla-optimal number of tokens.

"A quick(ish) Chinchilla check"

• gpjt/jax-with-mha-bias-larger-chinchilla-1 -- the slightly-larger model.
• gpjt/jax-with-mha-bias-larger-chinchilla-2 -- the slightly-smaller model.

I've also added links to the posts in question.

@misc{thomas2026sep-jax-models-on-hugging-face,
  author       = {Thomas, Giles},
  title        = {{Putting my JAX-trained models on the Hugging Face Hub}},
  year         = {2026},
  month        = sep,
  howpublished = {Blog post},
  url          = {https://www.gilesthomas.com/2026/09/jax-models-on-hugging-face},
}