Repositories list

• jax

  Public
  Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

  jax

  jax

  Python

  •

  Apache License 2.0

  •3.6k3.6k forks•36k36k stars•1.6k1.6k issues•744744 pull requests•Updated May 24, 2026May 24, 2026

• scaling-book

  Public
  Home for "How To Scale Your Model", a short blog-style textbook about scaling LLMs on TPUs

  rooflinejaxtpus+ 2

  rooflinejaxtpusllmsllm-inference

  HTML

  •

  MIT License

  •145145 forks•1k1k stars•44 issues•22 pull requests•Updated May 23, 2026May 23, 2026

• jax-tpu-embedding

  Public
  Python

  •

  Apache License 2.0

  •55 forks•3131 stars•00 issues•1010 pull requests•Updated May 23, 2026May 23, 2026

• jax-ai-stack

  Public
  Python

  •

  Apache License 2.0

  •4747 forks•296296 stars•99 issues•1515 pull requests•Updated May 11, 2026May 11, 2026

• jax-triton

  Public
  jax-triton contains integrations between JAX and OpenAI Triton

  Python

  •

  Apache License 2.0

  •5757 forks•460460 stars•1010 issues•2525 pull requests•Updated Apr 23, 2026Apr 23, 2026

• oryx

  Public
  Oryx is a library for probabilistic programming and deep learning built on top of Jax.

  pythonprobabilistic-programmingjax

  pythonprobabilistic-programmingjax

  Python

  •

  Apache License 2.0

  •1313 forks•316316 stars•1717 issues•33 pull requests•Updated Apr 20, 2026Apr 20, 2026

• ml_dtypes

  Public
  A stand-alone implementation of several NumPy dtype extensions used in machine learning.

  C++

  •

  Apache License 2.0

  •5757 forks•342342 stars•3131 issues•1818 pull requests•Updated Apr 15, 2026Apr 15, 2026

• jax-llm-examples

  Public
  Minimal yet performant LLM examples in pure JAX

  jaxllmllm-inference

  jaxllmllm-inference

  Python

  •

  Apache License 2.0

  •3636 forks•257257 stars•99 issues•00 pull requests•Updated Apr 10, 2026Apr 10, 2026

• bonsai

  Public
  Minimal, lightweight JAX implementations of popular models.

  pythonmachine-learningdeep-learning+ 2

  pythonmachine-learningdeep-learningjaxnnx

  Jupyter Notebook

  •

  Apache License 2.0

  •4444 forks•235235 stars•2424 issues•1414 pull requests•Updated Mar 27, 2026Mar 27, 2026

• bayeux

  Public
  State of the art inference for your bayesian models.

  Python

  •

  Apache License 2.0

  •1212 forks•241241 stars•88 issues•11 pull request•Updated Jan 14, 2026Jan 14, 2026

• jax-blog

  Public
  Apache License 2.0

  •22 forks•66 stars•00 issues•33 pull requests•Updated Jan 10, 2026Jan 10, 2026

• coix

  Public
  Inference Combinators in JAX

  Jupyter Notebook

  •

  Apache License 2.0

  •33 forks•5353 stars•99 issues•00 pull requests•Updated May 17, 2025May 17, 2025

• .allstar

  Public archive
  00 forks•00 stars•00 issues•00 pull requests•Updated Dec 6, 2022Dec 6, 2022

• australis

  Public archive
  C++

  •

  Apache License 2.0

  •55 forks•2828 stars•00 issues•00 pull requests•Updated Nov 18, 2022Nov 18, 2022

• .github

  Public archive
  00 forks•00 stars•00 issues•00 pull requests•Updated Oct 26, 2022Oct 26, 2022