Skip to content
@jax-ml

jax-ml

Pushing back the limits on numerical computing.

Pinned Loading

  1. jax Public

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

    Python 32.8k 3.1k

  2. jax-triton Public

    jax-triton contains integrations between JAX and OpenAI Triton

    Python 406 50

  3. ml_dtypes Public

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

    C++ 280 42

Repositories

Showing 10 of 14 repositories
  • jax Public

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

    Python 32,790 Apache-2.0 3,098 1,559 (5 issues need help) 580 Updated Jul 14, 2025
  • jax-tpu-embedding Public
    Python 18 Apache-2.0 3 0 23 Updated Jul 13, 2025
  • jax-llm-examples Public
    Python 132 Apache-2.0 17 2 1 Updated Jul 13, 2025
  • scaling-book Public

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

    HTML 428 MIT 55 0 3 Updated Jul 11, 2025
  • ml_dtypes Public

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

    C++ 280 Apache-2.0 42 27 8 Updated Jul 11, 2025
  • bonsai Public
    11 1 0 2 Updated Jul 11, 2025
  • oryx Public

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

    Python 269 Apache-2.0 10 17 (1 issue needs help) 2 Updated Jul 10, 2025
  • jax-ai-stack Public
    Python 196 Apache-2.0 28 7 14 Updated Jun 23, 2025
  • jax-triton Public

    jax-triton contains integrations between JAX and OpenAI Triton

    Python 406 Apache-2.0 50 8 19 Updated Jun 23, 2025
  • coix Public

    Inference Combinators in JAX

    Jupyter Notebook 50 Apache-2.0 3 9 0 Updated May 17, 2025
View all repositories

Most used topics

Loading…