JAX and Keras

Mentors: Student: Megan Andrews

As interest in large language models continues to grow, many open-source implementations exist across various frameworks—but accessible, modular, and well-documented implementations in Flax, built on JAX, remain limited. This lack of resources makes it more difficult for developers and researchers to understand, experiment with, and extend LLMs in the JAX ecosystem. This project addresses that gap by implementing one or more open LLM architectures in Flax. The goal is to provide reference models that are easy to read, modify, and build upon. The project will also include a comprehensive test suite to ensure reliability, and an interactive educational notebook that explains the architecture, shows how to load pretrained weights, and demonstrates how to run inference.

Mentors: Student: Bob Chen

Large Language Model (LLM) serving frameworks like vLLM must deliver high throughput and low latency to justify the substantial cost of GPU and TPU deployments. This project targets systematic benchmarking, profiling, and optimization of vLLM, one of the most widely adopted open-source LLM serving systems. The goal is to push its performance limits under realistic workloads and contribute meaningful improvements back to the community.

Mentors: Student: neel04

Current Frontier LLMs are really bad at assisting with writing JAX code due to the sparse resources spread across multiple different forums and platforms which are often inacessible/unscrapeable (such as discord) by conventional search engine indexers or web scrapers. Thus they usually do not make it into the final pretraining datasets, hurting frontier models' performance on those frameworks. I want to fine-tune OSS models on JAX code, where the data is manually scraped/collected by hand from diverse sources like YT talks, GitHub discussions, Books & Papers and even XLA HLO dumps to enhance the models' understanding of the compiler internals and improve debugging capabilities as well as better guide the developers when dealing with the internals. Another key focus would be incorporating commonly-used popular libraries like Equinox, Optax and Scalax as well because current SOTA LLMs are not familiar with these (relatively) niche but handy libraries and thus prove to be less helpful and prone to hallucinating components of the ecosystem that don't exist. This pretrained model will be available externally through an API or a CLI interface like Aider/Claude Code styled wherein it'd be accessible to everyone regardless of their development environment or IDE/Editor. I would also experiment with new ideas such as integrating type-annotations, Array and PyTree shape annotations and packaging other auxiliary information embedded in the AST to provide further context to the model apart from the raw code itself.

Mentors: Student: Dhiraj BM

VB-LoRA (Vector Bank Low-Rank Adaptation) is a parameter-efficient fine-tuning technique that extends LoRA (Low-Rank Adaptation) by introducing a shared vector bank and a differentiable top-k selection mechanism. This allows dense layers to adapt to new tasks by training a small set of parameters while keeping the original weights frozen. The goal of this project is to implement VB-LoRA in Keras, integrating it into layers—Dense, EinsumDense, BaseConv, Embedding, and Attention —and extending it to the backbones of model architectures in Keras Hub. The implementation will involve initializing VB-LoRA parameters, creating a vector bank and logits, computing low-rank matrices, and ensuring seamless integration with comprehensive testing and documentation.