Skip to content

@jax-ml jax-ml

Change the repository type filter

All

    Repositories list

    13 repositories

    • jax

      Public
      Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
      jax
      Python
      •
      Apache License 2.0
      •3k•32k•1.5k•488•Updated May 9, 2025May 9, 2025

    • jax-tpu-embedding

      Public
      Python
      •
      Apache License 2.0
      •3•17•0•15•Updated May 8, 2025May 8, 2025

    • jax-llm-examples

      Public
      Python
      •
      Apache License 2.0
      •14•109•1•1•Updated May 8, 2025May 8, 2025

    • oryx

      Public
      Oryx is a library for probabilistic programming and deep learning built on top of Jax.
      probabilistic-programmingjaxpython
      Python
      •
      Apache License 2.0
      •11•257•16•2•Updated May 8, 2025May 8, 2025

    • coix

      Public
      Inference Combinators in JAX
      Jupyter Notebook
      •
      Apache License 2.0
      •3•48•9•2•Updated May 7, 2025May 7, 2025

    • jax-ai-stack

      Public
      Python
      •
      Apache License 2.0
      •24•166•4•10•Updated May 7, 2025May 7, 2025

    • ml_dtypes

      Public
      A stand-alone implementation of several NumPy dtype extensions used in machine learning.
      C++
      •
      Apache License 2.0
      •38•262•23•5•Updated May 6, 2025May 6, 2025

    • jax-triton

      Public
      jax-triton contains integrations between JAX and OpenAI Triton
      Python
      •
      Apache License 2.0
      •46•393•8•18•Updated May 2, 2025May 2, 2025

    • bayeux

      Public
      State of the art inference for your bayesian models.
      Python
      •
      Apache License 2.0
      •13•213•6•1•Updated Apr 30, 2025Apr 30, 2025

    • scaling-book

      Public
      Home for "How To Scale Your Model", a short blog-style textbook about scaling LLMs on TPUs
      rooflinejaxtpusllmsllm-inference
      HTML
      •
      MIT License
      •42•259•0•1•Updated Apr 24, 2025Apr 24, 2025

    • .allstar

      Public archive
      0•0•0•0•Updated Dec 6, 2022Dec 6, 2022

    • australis

      Public archive
      C++
      •
      Apache License 2.0
      •5•29•0•0•Updated Nov 18, 2022Nov 18, 2022

    • .github

      Public archive
      0•0•0•0•Updated Oct 26, 2022Oct 26, 2022