mcx by rlouf

Express & compile probabilistic programs for performant inference on CPU & GPU. Powered by JAX.

created at Jan. 22, 2020, 8:38 a.m.

17 +0

324 +0

17 +0

tinygp by dfm

The tiniest of Gaussian Process libraries

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

9 +0

280 +2

22 +0

jax-flows by ChrisWaites

Normalizing Flows in JAX 🌊

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

8 +0

263 +0

17 +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

249 +0

41 +0

QDax by adaptive-intelligent-robotics

Accelerated Quality-Diversity

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

6 +0

243 +0

35 +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

223 +1

24 +0

ott by google-research

None

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

15 +0

215 +0

18 +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.

10 +0

202 +1

15 +0

jaxlie by brentyi

Rigid transforms + Lie groups in JAX

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

4 +0

202 +1

14 +0

JAX-Toolbox by NVIDIA

JAX-Toolbox

created at Feb. 13, 2023, 2:39 p.m.

23 +0

177 +6

34 +1

tree-math by google

Mathematical operations for JAX pytrees

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

10 -1

167 +0

7 +0

coax by coax-dev

Modular framework for Reinforcement Learning in python

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

8 +0

161 +0

18 +0

parallax by srush

None

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

6 +0

157 +0

4 +0

EasyDeL by erfanzar

EasyDeL is an OpenSource Library to make your training faster and more Optimized With cool Options for training and serving Both in Python And Mojo🔥

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

9 +0

157 +5

17 +0

jax_cosmo by DifferentiableUniverseInitiative

A differentiable cosmology library in JAX

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

18 +0

157 +0

30 +0

XLB by Autodesk

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

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

10 +0

146 +3

15 +0

praxis by google

None

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

8 +0

144 +3

39 +3

jax-models by DarshanDeshpande

Unofficial JAX implementations of deep learning research papers

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

5 +0

140 +0

9 +0

jaxns by Joshuaalbert

Probabilistic Programming and Nested sampling in JAX

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

5 +0

125 +4

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 +1

124 +1

21 +1