jax-ml

Top repositories by stars

• jax-ml/jax(on leaderboard)

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

  Python34,897
• jax-ml/scaling-book

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

  HTML853
• jax-ml/jax-triton

  jax-triton contains integrations between JAX and OpenAI Triton

  Python440
• jax-ml/ml_dtypes

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

  C++330
• jax-ml/oryx

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

  Python298
• jax-ml/jax-ai-stack

  Python268
• jax-ml/jax-llm-examples

  Minimal yet performant LLM examples in pure JAX

  Python242
• jax-ml/bayeux

  State of the art inference for your bayesian models.

  Python232
• jax-ml/bonsai

  Minimal, lightweight JAX implementations of popular models.

  Jupyter Notebook198
• jax-ml/coix

  Inference Combinators in JAX

  Jupyter Notebook52
• jax-ml/australis

  C++27
• jax-ml/jax-tpu-embedding

  Python26
• jax-ml/jax-blog

  6
• jax-ml/.github

  0
• jax-ml/.allstar

  0