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 32k 3k

  2. jax-triton Public

    jax-triton contains integrations between JAX and OpenAI Triton

    Python 390 45

  3. ml_dtypes Public

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

    C++ 256 36

Repositories

Showing 10 of 13 repositories
  • jax Public

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

    Python 31,958 Apache-2.0 2,980 1,485 (5 issues need help) 475 Updated Apr 17, 2025
  • oryx Public

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

    Python 254 Apache-2.0 11 15 (1 issue needs help) 3 Updated Apr 17, 2025
  • scaling-book Public

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

    HTML 240 MIT 36 0 3 Updated Apr 16, 2025
  • jax-tpu-embedding Public
    Python 17 Apache-2.0 3 0 11 Updated Apr 16, 2025
  • jax-ai-stack Public
    Python 154 Apache-2.0 25 3 10 Updated Apr 16, 2025
  • bayeux Public

    State of the art inference for your bayesian models.

    Python 209 Apache-2.0 11 7 0 Updated Apr 11, 2025
  • jax-triton Public

    jax-triton contains integrations between JAX and OpenAI Triton

    Python 390 Apache-2.0 45 8 18 Updated Apr 10, 2025
  • jax-llm-examples Public
    Python 98 Apache-2.0 11 1 1 Updated Apr 9, 2025
  • ml_dtypes Public

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

    C++ 256 Apache-2.0 36 22 5 Updated Mar 31, 2025
  • coix Public

    Inference Combinators in JAX

    Jupyter Notebook 47 Apache-2.0 2 9 0 Updated Oct 24, 2024
View all repositories

Most used topics

Loading…