Skip to content

@jax-ml jax-ml

Change the repository type filter

All

    Repositories list

    15 repositories

    • jax

      Public
      Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
      jax
      jax
      Python
      •
      Apache License 2.0
      •3.5k•35k•1.6k•660•Updated Apr 27, 2026Apr 27, 2026

    • jax-tpu-embedding

      Public
      Python
      •
      Apache License 2.0
      •5•31•0•7•Updated Apr 24, 2026Apr 24, 2026

    • jax-triton

      Public
      jax-triton contains integrations between JAX and OpenAI Triton
      Python
      •
      Apache License 2.0
      •57•447•10•21•Updated Apr 23, 2026Apr 23, 2026

    • jax-ai-stack

      Public
      Python
      •
      Apache License 2.0
      •47•291•9•16•Updated Apr 20, 2026Apr 20, 2026

    • oryx

      Public
      Oryx is a library for probabilistic programming and deep learning built on top of Jax.
      pythonprobabilistic-programmingjax
      pythonprobabilistic-programmingjax
      Python
      •
      Apache License 2.0
      •13•313•17•3•Updated Apr 20, 2026Apr 20, 2026

    • ml_dtypes

      Public
      A stand-alone implementation of several NumPy dtype extensions used in machine learning.
      C++
      •
      Apache License 2.0
      •56•340•30•16•Updated Apr 15, 2026Apr 15, 2026

    • jax-llm-examples

      Public
      Minimal yet performant LLM examples in pure JAX
      jaxllmllm-inference
      jaxllmllm-inference
      Python
      •
      Apache License 2.0
      •34•251•9•0•Updated Apr 10, 2026Apr 10, 2026

    • bonsai

      Public
      Minimal, lightweight JAX implementations of popular models.
      pythonmachine-learningdeep-learning
      pythonmachine-learningdeep-learningjaxnnx
      Jupyter Notebook
      •
      Apache License 2.0
      •43•234•23•13•Updated Mar 27, 2026Mar 27, 2026

    • scaling-book

      Public
      Home for "How To Scale Your Model", a short blog-style textbook about scaling LLMs on TPUs
      rooflinejaxtpus
      rooflinejaxtpusllmsllm-inference
      HTML
      •
      MIT License
      •133•928•4•1•Updated Mar 15, 2026Mar 15, 2026

    • bayeux

      Public
      State of the art inference for your bayesian models.
      Python
      •
      Apache License 2.0
      •12•240•8•1•Updated Jan 14, 2026Jan 14, 2026

    • jax-blog

      Public
      Apache License 2.0
      •2•6•0•3•Updated Jan 10, 2026Jan 10, 2026

    • coix

      Public
      Inference Combinators in JAX
      Jupyter Notebook
      •
      Apache License 2.0
      •3•53•9•0•Updated May 17, 2025May 17, 2025

    • .allstar

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

    • australis

      Public archive
      C++
      •
      Apache License 2.0
      •4•28•0•0•Updated Nov 18, 2022Nov 18, 2022

    • .github

      Public archive
      0•0•0•0•Updated Oct 26, 2022Oct 26, 2022
    ProTip! When viewing an organization's repositories, you can use the props. filter to filter by custom property.