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.9k3.9k forks in jax-ml/jax•37k37k stars in jax-ml/jax•1.7k1.7k issues in jax-ml/jax•854854 pull requests in jax-ml/jax•Updated Oct 28, 2026Oct 28, 2026

• jax-tpu-embedding

  Public
  High-performance large-scale embedding acceleration for JAX on Google TPU SparseCores.

  deep-learningembeddingstpu+ 4

  deep-learningembeddingstpujaxxlasparse-computationjax-sc

  Python

  •

  Apache License 2.0

  •66 forks in jax-ml/jax-tpu-embedding•3838 stars in jax-ml/jax-tpu-embedding•00 issues in jax-ml/jax-tpu-embedding•99 pull requests in jax-ml/jax-tpu-embedding•Updated Oct 28, 2026Oct 28, 2026

• ml_dtypes

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

  C++

  •

  Apache License 2.0

  •6464 forks in jax-ml/ml_dtypes•363363 stars in jax-ml/ml_dtypes•3030 issues in jax-ml/ml_dtypes•1616 pull requests in jax-ml/ml_dtypes•Updated Oct 27, 2026Oct 27, 2026

• bayeux

  Public
  State of the art inference for your bayesian models.

  Python

  •

  Apache License 2.0

  •1515 forks in jax-ml/bayeux•249249 stars in jax-ml/bayeux•88 issues in jax-ml/bayeux•44 pull requests in jax-ml/bayeux•Updated Oct 27, 2026Oct 27, 2026

• jax-triton

  Public
  jax-triton contains integrations between JAX and OpenAI Triton

  Python

  •

  Apache License 2.0

  •6161 forks in jax-ml/jax-triton•473473 stars in jax-ml/jax-triton•44 issues in jax-ml/jax-triton•77 pull requests in jax-ml/jax-triton•Updated Oct 20, 2026Oct 20, 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

  SCSS

  •

  MIT License

  •214214 forks in jax-ml/scaling-book•1.5k1.5k stars in jax-ml/scaling-book•44 issues in jax-ml/scaling-book•99 pull requests in jax-ml/scaling-book•Updated Oct 19, 2026Oct 19, 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 in jax-ml/oryx•326326 stars in jax-ml/oryx•1717 issues in jax-ml/oryx•33 pull requests in jax-ml/oryx•Updated Oct 14, 2026Oct 14, 2026

• jax-ai-stack

  Public
  Python

  •

  Apache License 2.0

  •5151 forks in jax-ml/jax-ai-stack•314314 stars in jax-ml/jax-ai-stack•1313 issues in jax-ml/jax-ai-stack•2525 pull requests in jax-ml/jax-ai-stack•Updated Oct 13, 2026Oct 13, 2026

• jax-llm-examples

  Public
  Minimal yet performant LLM examples in pure JAX

  jaxllmllm-inference

  jaxllmllm-inference

  Python

  •

  Apache License 2.0

  •3838 forks in jax-ml/jax-llm-examples•286286 stars in jax-ml/jax-llm-examples•99 issues in jax-ml/jax-llm-examples•00 pull requests in jax-ml/jax-llm-examples•Updated Aug 23, 2026Aug 23, 2026

• bonsai

  Public
  Minimal, lightweight JAX implementations of popular models.

  pythonmachine-learningdeep-learning+ 2

  pythonmachine-learningdeep-learningjaxnnx

  Python

  •

  Apache License 2.0

  •4646 forks in jax-ml/bonsai•246246 stars in jax-ml/bonsai•2424 issues in jax-ml/bonsai•1414 pull requests in jax-ml/bonsai•Updated May 29, 2026May 29, 2026

• jax-blog

  Public
  Apache License 2.0

  •22 forks in jax-ml/jax-blog•66 stars in jax-ml/jax-blog•00 issues in jax-ml/jax-blog•33 pull requests in jax-ml/jax-blog•Updated Jan 10, 2026Jan 10, 2026

• coix

  Public
  Inference Combinators in JAX

  Jupyter Notebook

  •

  Apache License 2.0

  •33 forks in jax-ml/coix•5454 stars in jax-ml/coix•99 issues in jax-ml/coix•00 pull requests in jax-ml/coix•Updated May 17, 2025May 17, 2025

• .allstar

  Public archive
  00 forks in jax-ml/.allstar•00 stars in jax-ml/.allstar•00 issues in jax-ml/.allstar•00 pull requests in jax-ml/.allstar•Updated Dec 6, 2022Dec 6, 2022

• australis

  Public archive
  C++

  •

  Apache License 2.0

  •66 forks in jax-ml/australis•2828 stars in jax-ml/australis•00 issues in jax-ml/australis•00 pull requests in jax-ml/australis•Updated Nov 18, 2022Nov 18, 2022

• .github

  Public archive
  00 forks in jax-ml/.github•00 stars in jax-ml/.github•00 issues in jax-ml/.github•00 pull requests in jax-ml/.github•Updated Oct 26, 2022Oct 26, 2022