Skip to content
@jax-ml

jax-ml

Pushing back the limits on numerical computing.

Pinned Loading

  1. jax jax Public

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

    Python 36.4k 3.8k

  2. jax-llm-examples jax-llm-examples Public

    Minimal yet performant LLM examples in pure JAX

    Python 285 38

  3. jax-triton jax-triton Public

    jax-triton contains integrations between JAX and OpenAI Triton

    Python 471 60

  4. scaling-book scaling-book Public

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

    SCSS 1.5k 212

  5. ml_dtypes ml_dtypes Public

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

    C++ 361 60

Repositories

Showing 10 of 15 repositories
  • jax Public

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

    jax-ml/jax's past year of commit activity
    Python 36,370 Apache-2.0 3,828 1,777 (3 issues need help) 853 Updated Oct 2, 2026
  • jax-tpu-embedding Public

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

    jax-ml/jax-tpu-embedding's past year of commit activity
    Python 37 Apache-2.0 5 0 5 Updated Oct 2, 2026
  • oryx Public

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

    jax-ml/oryx's past year of commit activity
    Python 328 Apache-2.0 13 17 (1 issue needs help) 3 Updated Sep 29, 2026
  • jax-triton Public

    jax-triton contains integrations between JAX and OpenAI Triton

    jax-ml/jax-triton's past year of commit activity
    Python 471 Apache-2.0 60 4 7 Updated Sep 28, 2026
  • ml_dtypes Public

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

    jax-ml/ml_dtypes's past year of commit activity
    C++ 361 Apache-2.0 60 29 12 Updated Sep 24, 2026
  • scaling-book Public

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

    jax-ml/scaling-book's past year of commit activity
    SCSS 1,458 MIT 212 3 10 Updated Sep 22, 2026
  • jax-ai-stack Public
    jax-ml/jax-ai-stack's past year of commit activity
    Python 311 Apache-2.0 50 13 23 Updated Sep 21, 2026
  • jax-llm-examples Public

    Minimal yet performant LLM examples in pure JAX

    jax-ml/jax-llm-examples's past year of commit activity
    Python 285 Apache-2.0 38 9 0 Updated Aug 23, 2026
  • bonsai Public

    Minimal, lightweight JAX implementations of popular models.

    jax-ml/bonsai's past year of commit activity
    Python 245 Apache-2.0 47 24 (2 issues need help) 14 Updated May 29, 2026
  • bayeux Public

    State of the art inference for your bayesian models.

    jax-ml/bayeux's past year of commit activity
    Python 248 Apache-2.0 15 8 3 Updated Jan 14, 2026
View all repositories

Top languages

Loading…