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 34.9k 3.4k

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

    Minimal yet performant LLM examples in pure JAX

    Python 242 31

  3. jax-triton jax-triton Public

    jax-triton contains integrations between JAX and OpenAI Triton

    Python 440 54

  4. scaling-book scaling-book Public

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

    HTML 848 122

  5. ml_dtypes ml_dtypes Public

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

    C++ 330 55

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 34,866 Apache-2.0 3,420 1,672 (5 issues need help) 793 Updated Feb 15, 2026
  • jax-tpu-embedding Public
    jax-ml/jax-tpu-embedding’s past year of commit activity
    Python 26 Apache-2.0 5 0 9 Updated Feb 14, 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
    HTML 848 MIT 122 3 1 Updated Feb 14, 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++ 330 Apache-2.0 55 30 10 Updated Feb 12, 2026
  • jax-ai-stack Public
    jax-ml/jax-ai-stack’s past year of commit activity
    Python 267 Apache-2.0 46 9 20 Updated Feb 9, 2026
  • jax-triton Public

    jax-triton contains integrations between JAX and OpenAI Triton

    jax-ml/jax-triton’s past year of commit activity
    Python 440 Apache-2.0 54 10 21 Updated Feb 9, 2026
  • bonsai Public

    Minimal, lightweight JAX implementations of popular models.

    jax-ml/bonsai’s past year of commit activity
    Jupyter Notebook 196 Apache-2.0 42 22 (2 issues need help) 18 Updated Feb 6, 2026
  • bayeux Public

    State of the art inference for your bayesian models.

    jax-ml/bayeux’s past year of commit activity
    Python 232 Apache-2.0 12 8 1 Updated Jan 14, 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 242 Apache-2.0 31 8 1 Updated Jan 14, 2026
  • jax-blog Public
    jax-ml/jax-blog’s past year of commit activity
    6 Apache-2.0 2 0 3 Updated Jan 10, 2026
View all repositories