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 31.4k 2.9k

  2. jax-triton jax-triton Public

    jax-triton contains integrations between JAX and OpenAI Triton

    Python 379 42

  3. ml_dtypes ml_dtypes Public

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

    C++ 253 30

Repositories

Showing 10 of 13 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 31,414 Apache-2.0 2,896 1,499 (6 issues need help) 445 Updated Feb 25, 2025
  • jax-tpu-embedding Public
    jax-ml/jax-tpu-embedding’s past year of commit activity
    Python 14 Apache-2.0 3 0 9 Updated Feb 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++ 253 Apache-2.0 30 20 7 Updated Feb 24, 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 191 MIT 28 1 2 Updated Feb 24, 2025
  • examples Public
    jax-ml/examples’s past year of commit activity
    0 Apache-2.0 0 0 0 Updated Feb 24, 2025
  • jax-ai-stack Public
    jax-ml/jax-ai-stack’s past year of commit activity
    Python 129 Apache-2.0 22 4 7 Updated Feb 24, 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 246 Apache-2.0 10 15 (1 issue needs help) 4 Updated Feb 13, 2025
  • jax-triton Public

    jax-triton contains integrations between JAX and OpenAI Triton

    jax-ml/jax-triton’s past year of commit activity
    Python 379 Apache-2.0 42 37 (1 issue needs help) 19 Updated Jan 23, 2025
  • bayeux Public

    State of the art inference for your bayesian models.

    jax-ml/bayeux’s past year of commit activity
    Python 200 Apache-2.0 9 4 0 Updated Dec 10, 2024
  • coix Public

    Inference Combinators in JAX

    jax-ml/coix’s past year of commit activity
    Jupyter Notebook 46 Apache-2.0 2 9 0 Updated Oct 24, 2024
View all repositories

Most used topics

Loading…