jax-ml
OrganizationPushing back the limits on numerical computing.
Public repos
15
Public gists
0
Member since
Dec 3, 2019
On the leaderboard
| Rank | Repository | Stars |
|---|---|---|
| 752 | jax-ml/jax | 35,302 |
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 - 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 - C++27
- Python26