JAX
216 items from n2cholas/awesome-jax ★2,167
-
gymnax github.com
Reinforcement Learning Environments with the well-known gym API.
-
mip-NeRF github.com
Official implementation of Mip-NeRF: A Multiscale Representation for Anti-Aliasing Neural Radiance Fields.
-
Fortuna github.com
AWS library for Uncertainty Quantification in Deep Learning.
-
Jumanji github.com
A Suite of Industry-Driven Hardware-Accelerated RL Environments written in JAX.
-
evosax github.com
JAX-Based Evolution Strategies
-
Get started with JAX by Aleksa Gordić github.com
A series of notebooks and videos going from zero JAX knowledge to building neural networks in Haiku.
-
Objax github.com
Has an object oriented design similar to PyTorch.
-
OTT-JAX github.com
Optimal transport tools in JAX.
-
JAX RL github.com
Implementations of reinforcement learning algorithms.
-
Levanter github.com
Legible, Scalable, Reproducible Foundation Models with Named Tensors and JAX.
-
BrainPy github.com
Brain Dynamics Programming in Python.
-
NetKet github.com
Machine Learning toolbox for Quantum Physics.
-
GPJax github.com
Gaussian processes in JAX.
-
Distrax github.com
Reimplementation of TensorFlow Probability, containing probability distributions and bijectors.
-
Pgx github.com
Vectorized board game environments for RL with an AlphaZero example.
-
Optimistix github.com
Root finding, minimisation, fixed points, and least squares.
-
Pax github.com
A Jax-based machine learning framework for training large scale models.
-
mpi4jax github.com
Combine MPI operations with your Jax code on CPUs and GPUs.
-
XLB github.com
A Differentiable Massively Parallel Lattice Boltzmann Library in Python for Physics-Based Machine Learning.
-
econpizza github.com
Solve macroeconomic models with hetereogeneous agents using JAX.
-
PIX github.com
PIX is an image processing library in JAX, for JAX.
-
JAX Toolbox github.com
Nightly CI and optimized examples for JAX on NVIDIA GPUs using libraries such as T5x, Paxml, and Transformer Engine.
-
Extending JAX with custom C++ and CUDA code by Dan Foreman-Mackey github.com
Tutorial demonstrating the infrastructure required to provide custom ops in JAX.
-
EasyDeL github.com
EasyDeL 🔮 is an OpenSource Library to make your training faster and more Optimized With cool Options for training and serving (Llama, MPT, Mixtral, Falcon, etc) in JAX
-
QDax github.com
Quality Diversity optimization in Jax.
-
FDTDX github.com
Finite-Difference Time-Domain Electromagnetic Simulations in JAX
-
tinygp github.com
The tiniest of Gaussian process libraries in JAX.
-
jaxlie github.com
Lie theory library for rigid body transformations and optimization.
-
mcx github.com
Express & compile probabilistic programs for performant inference.
-
KFAC-JAX github.com
Second Order Optimization with Approximate Curvature for NNs.
-
SPU github.com
A domain-specific compiler and runtime suite to run JAX code with MPC(Secure Multi-Party Computation).
-
dynamiqs github.com
High-performance and differentiable simulations of quantum systems with JAX.
-
Equivariant MLP github.com
Construct equivariant neural network layers.
-
FedJAX github.com
Federated learning in JAX, built on Optax and Haiku.
-
flaxmodels github.com
Pretrained models for Jax/Flax.
-
jaxns github.com
Nested sampling in JAX.
-
torchax github.com
torchax is a library for Jax to interoperate with model code written in PyTorch.
-
jax-cosmo github.com
Differentiable cosmology library.
-
flowjax github.com
Distributions and normalizing flows built as equinox modules.
-
jwave github.com
A library for differentiable acoustic simulations
-
Optimal Transport Tools github.com
Toolbox that bundles utilities to solve optimal transport problems.
-
tree-math github.com
Convert functions that operate on arrays into functions that operate on PyTrees.
-
NesT github.com
Official implementation of Aggregating Nested Transformers.
-
Praxis github.com
The layer library for Pax with a goal to be usable by other JAX-based ML projects.
-
Coax github.com
Turn RL papers into code, the easy way.
-
NAVIX github.com
A reimplementation of MiniGrid, a Reinforcement Learning environment, in JAX
-
kvax github.com
A FlashAttention implementation for JAX with support for efficient document mask computation and context parallelism.
-
SCICO github.com
Scientific computational imaging in JAX.
-
jax-models github.com
Implementations of research papers originally without code or code written with frameworks other than JAX.
-
Parallax github.com
Immutable Torch Modules for JAX.
- next page of items loading…
