tinygp by dfm

The tiniest of Gaussian Process libraries

created at Feb. 14, 2021, 3:51 p.m.

10 +1

298 +1

24 +0

jax-flows by ChrisWaites

Normalizing Flows in JAX 🌊

created at March 23, 2020, 2:42 a.m.

8 +0

274 +0

19 +0

QDax by adaptive-intelligent-robotics

Accelerated Quality-Diversity

created at Feb. 11, 2022, 3:48 p.m.

5 +0

268 +0

46 +0

fedjax by google

FedJAX is a JAX-based open source library for Federated Learning simulations that emphasizes ease-of-use in research.

created at Dec. 22, 2020, 2:44 p.m.

11 +0

254 +0

41 +0

kfac-jax by deepmind

Second Order Optimization and Curvature Estimation with K-FAC in JAX.

created at March 18, 2022, 10:19 a.m.

11 +0

251 +1

23 +0

flaxmodels by matthias-wright

Pretrained deep learning models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet, etc.

created at May 16, 2021, 1:34 p.m.

6 +0

238 +0

24 +0

XLB by Autodesk

XLB: Accelerated Lattice Boltzmann (XLB) for Physics-based ML

created at May 4, 2023, 5:36 p.m.

12 +0

238 +2

25 +0

jaxlie by brentyi

Rigid transforms + Lie groups in JAX

created at Nov. 28, 2020, 6:18 a.m.

7 +1

234 +0

15 +0

ott by google-research

None

created at Jan. 18, 2021, 3:03 p.m.

15 +0

213 +0

18 +0

EasyDeL by erfanzar

Accelerate, Optimize performance with streamlined training and serving options with JAX.

created at April 9, 2023, 5:18 p.m.

13 +0

208 +0

25 +0

tree-math by google

Mathematical operations for JAX pytrees

created at Dec. 18, 2021, 1:56 a.m.

11 +0

190 +1

8 +1

jax_cosmo by DifferentiableUniverseInitiative

A differentiable cosmology library in JAX

created at May 1, 2020, 5:47 p.m.

17 +0

178 -1

37 +0

praxis by google

None

created at June 14, 2022, 4:04 p.m.

8 +0

177 +0

43 +0

dynamiqs by dynamiqs

High-performance quantum systems simulation with JAX (GPU-accelerated & differentiable solvers).

created at Feb. 5, 2023, 5:04 p.m.

6 +0

175 +4

18 +1

coax by coax-dev

Modular framework for Reinforcement Learning in python

created at March 11, 2021, 12:56 a.m.

8 +0

169 +1

17 +0

parallax by srush

None

created at May 19, 2020, 1:49 a.m.

6 +0

155 +0

4 +0

jaxns by Joshuaalbert

Probabilistic Programming and Nested sampling in JAX

created at July 31, 2020, 6 p.m.

5 +0

153 +3

10 +0

jax-models by DarshanDeshpande

Unofficial JAX implementations of deep learning research papers

created at Jan. 8, 2022, 1:40 p.m.

5 +0

151 +0

8 +0

jwave by ucl-bug

A JAX-based research framework for differentiable and parallelizable acoustic simulations, on CPU, GPUs and TPUs

created at April 6, 2021, 10:22 a.m.

8 +0

143 +0

21 +0

lorax by davisyoshida

LoRA for arbitrary JAX models and functions

created at April 21, 2023, 11:50 a.m.

4 +0

134 +1

5 +0