Top AI Repos — open-source AI, indexed and scored
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
JAX library for training sub-4B foundation models for edge
| Date | Stars |
|---|---|
| 2026-07-24 | 306 |
| 2026-07-25 | 306 |
| 2026-07-28 | 306 |
| 2026-07-30 | 306 |
| 2026-07-31 | 306 |
| 2026-08-06 | 306 |
Today
— stars today
This week
— stars this week
This month
— stars this month
Momentum
0.0
growth rate 0.00%/day
<p align="center"> <img src="assets/logo.jpg" alt="Alt text"/> </p> # A Jax-based library for designing and training transformer models from scratch.  [](https://nanodl.readthedocs.io) [](https://discord.gg/3u9vumJEmz) [](https://www.linkedin.com//company/80434055) [](https://twitter.com/hmunachii) Author: [Henry Ndubuaku](https://www.linkedin.com/in/henry-ndubuaku-7b6350b8/) (Discord & Docs badges are clickable) N/B: Codes are implemented pedagogically at the expense of repetition. Each model is purposefully contained in a file without inter-file dependencies. ## Overview Developing and training transformer-based models is typically resource-intensive and time-consuming and AI/ML experts frequently need to build smaller-scale versions of these models for specific problems. Jax, a low-resource yet powerful framework, accelerates the development of neural networks and abstracts distributed training, but existing resources for transformer development in Jax are limited. NanoDL addresses this challenge with the following features: - A wide array of blocks and layers, facilitating the creation of customised transformer models from scratch. - An extensive selection of models like Gemma, LlaMa3, Mistral, GPT3, GPT4 (inferred), T5, Whisper, ViT, Mixers, CLIP etc. - Data-parallel distributed trainers models on multiple GPUs or TPUs, without the need for manual training loops. - Dataloaders, making the process of data handling for Jax/Flax more straightforward and effective. - Layers not found in Flax/Jax, such as RoPE, GQA, MQA, and SWin attention, allowing for more flexible model development. - GPU/TPU-accelerated classical ML models like PCA, KMeans, Regression, Gaussian Processes etc. - True random number generators in Jax which do not need the verbose code. - A range of advanced algorithms for NLP and computer vision tasks, such as Gaussian Blur, BLEU, Tokenizer etc. - Each model is contained in a single file with no external dependencies, so the source code can also be easily used. - True random number generators in Jax which do not need the verbose code (examples shown in next sections). There are experimental and/or unfinished features (like MAMBA, KAN, BitNet, GAT and RLHF) in the repo which are not yet available via the package, but can be copied from this repo. Feedback on any of our discussion, issue and pull request threads are welcomed! Please report any feature requests, issues, questions or concerns in the [Discord](https://discord.gg/3u9vumJEmz), or just let us know what you're working on! ## Quick install You will need Python 3.9 or later, and working [JAX](https://github.com/google/jax/blob/main/README.md) installation, [FLAX](https://github.com/google/flax/blob/main/README.md) installation, [OPTAX](https://github.com/google-deepmind/optax/blob/main/README.md) installation (with GPU support for running training, without can only support creations). Models can be designed and tested on CPUs but trainers are all Distributed Data-Parallel which would require a GPU with 1 to N GPUS/TPUS. For CPU-only version of JAX: ``` pip install --upgrade pip # To support manylinux2010 wheels. pip install jax flax optax ``` Then, install nanodl from PyPi: ``` pip install nanodl ``` ## What does nanodl look like? We provide various example usages of the nanodl API. ```py import jax import nanodl import jax.numpy as jnp from nanodl import ArrayDataset, DataLoader from nanodl import GPT4, GPTDataParallelTrainer # Preparing your dataset batch_size = 8 max_length
Excerpt of 11,565 characters
Read on GitHubWould you bet a product on this? Bounded 0–100 and slow moving.
matched fp:0d474f8adcd119db, topic:transformer, topic:gpt, topic:llama
matched fp:0d474f8adcd119db, topic:deep-learning, topic:jax, topic:distributed-training