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.2k 3.3k

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

    Minimal yet performant LLM examples in pure JAX

    Python 204 27

  3. jax-triton jax-triton Public

    jax-triton contains integrations between JAX and OpenAI Triton

    Python 436 53

  4. scaling-book scaling-book Public

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

    HTML 710 104

  5. ml_dtypes ml_dtypes Public

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

    C++ 313 50

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,166 Apache-2.0 3,281 1,622 (5 issues need help) 600 Updated Dec 1, 2025
  • jax-tpu-embedding Public
    jax-ml/jax-tpu-embedding’s past year of commit activity
    Python 25 Apache-2.0 4 0 7 Updated Dec 1, 2025
  • jax-ai-stack Public
    jax-ml/jax-ai-stack’s past year of commit activity
    Python 246 Apache-2.0 42 8 18 Updated Dec 1, 2025
  • 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 710 MIT 104 4 1 Updated Nov 29, 2025
  • bonsai Public

    Minimal, lightweight JAX implementations of popular models.

    jax-ml/bonsai’s past year of commit activity
    Jupyter Notebook 166 Apache-2.0 28 14 (2 issues need help) 3 Updated Nov 25, 2025
  • 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++ 313 Apache-2.0 50 29 6 Updated Nov 25, 2025
  • jax-triton Public

    jax-triton contains integrations between JAX and OpenAI Triton

    jax-ml/jax-triton’s past year of commit activity
    Python 436 Apache-2.0 53 10 25 Updated Nov 24, 2025
  • jax-blog Public
    jax-ml/jax-blog’s past year of commit activity
    3 Apache-2.0 1 0 1 Updated Nov 21, 2025
  • 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 285 Apache-2.0 12 17 (1 issue needs help) 3 Updated Oct 7, 2025
  • jax-llm-examples Public

    Minimal yet performant LLM examples in pure JAX

    jax-ml/jax-llm-examples’s past year of commit activity
    Python 204 Apache-2.0 27 8 5 Updated Sep 23, 2025
View all repositories