GSoC 2025: "Implement Open LLM Models with JAX and Flax" Project

[akshayatam]

Hi everyone,

I'm currently preparing a proposal for the “Implement Open LLM Models with JAX and Flax” project for Google Summer of Code 2025, and I’m very excited about the opportunity to contribute to the JAX and Flax ecosystems. I noticed the email list on the page is unresponsive at the time of posting.

I have a background in machine learning, and recently completed a project where I built a DC-GAN from scratch using JAX to deepen my understanding of the framework. I’ve also worked on hate speech detection using DistilBERT, as well as a machine translation system using RetNet (a novel alternative to multi-head attention).

This project aligns closely with my goal of developing expertise in model implementation and optimization using JAX and Flax, and contributing reference models and educational material that helps others build with these tools.

As I finalize my proposal, I would really appreciate feedback or clarification on a few points:

• Model Selection: Are there specific open-source LLMs you'd recommend starting with? (like the LLaMA family or the Mistral family)
• Performance Goals: Is optimizing for TPU deployment a strong focus, or would a GPU baseline be sufficient initially?
• Preferred APIs: I noticed flax.nnx was used for the miniGPT tutorial on the project description page. Should the project use flax.linen, flax.nnx, or both depending on the use case?
• Evaluation Expectations: Is it expected that the implementation include a benchmark comparison or performance metrics, or is functionality and clarity of code the primary goal?

I’m committed to making meaningful contributions during and beyond GSoC and would love to align my proposal with what would be most useful to the community.