#jax (30 Repositories)
Ranked open-source repositories tagged with #jax, scored by pull request acceptance likelihood and maintainer engagement velocity.
69.3%
20.4h
30 repositories tagged #jax
GalacticDynamics/unxt
Unitful Quantities in JAX
nstarman/quax
Multiple dispatch over abstract array types in JAX.
kubeflow/trainer
Distributed AI Model Training and LLM Fine-Tuning on Kubernetes
ekzhang/jax-js
JAX in JavaScript – ML library for the web, running on WebGPU & Wasm
PennyLaneAI/catalyst
A JIT compiler for hybrid quantum programs in PennyLane
google/orbax
Orbax provides common checkpointing and persistence utilities for JAX users
keras-team/keras
Deep Learning for humans
iree-org/iree
A retargetable MLIR-based machine learning compiler and runtime toolkit.
NVIDIA/TransformerEngine
A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit and 4-bit floating point (FP8 and FP4) precision on Hopper, Ada and Blackwell GPUs, to provide better performance with lower memory utilization in both training and inference.
lanl/scico
Scientific Computational Imaging COde
pyro-ppl/numpyro
Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.
jax-ml/jax
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
keras-team/keras-hub
Pretrained model hub for Keras 3.
tillahoffmann/jax-mps
A JAX backend for Apple Metal Performance Shaders (MPS), enabling GPU-accelerated JAX computations on Apple Silicon.
Joshuaalbert/jaxns
Probabilistic Programming and Nested sampling in JAX
thomaspinder/GPJax
Gaussian processes in JAX and Equinox.
AI-Hypercomputer/maxtext
A simple, performant, and scalable Jax LLM!
HIPS/autograd
Efficiently computes derivatives of NumPy code.
team-ocean/veros
The versatile ocean simulator, in pure Python, powered by JAX.
ComputationalPsychiatry/pyhgf
PyHGF: A neural network library for predictive coding
google-research/scenic
Scenic: A Jax Library for Computer Vision Research and Beyond
instadeepai/Mava
🦁 A research-friendly codebase for fast experimentation of multi-agent reinforcement learning in JAX
kubeflow/sdk
Universal Python SDK to run AI workloads on Kubernetes
bardhh/cbfkit
JAX toolbox for safe robotics control: Control Barrier Function (CBF) safety filters, MPPI planning, and a drop-in safe-RL layer.
CogniPilot/rumoca
a rust-native Modelica compiler for fast, composable, scientific computing
patrick-kidger/jaxtyping
Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/
QuantClimate/GPJax
Gaussian processes in JAX and Equinox.
jaketae/koclip
KoCLIP: Korean port of OpenAI CLIP, in Flax
flatironinstitute/jaxmg
JAXMg: A multi-GPU linear solver in JAX
magenta/magenta-realtime
Magenta RealTime 2: An Open-Weights Live Music Model