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

24 +0

deltapv by romanodev

A photovoltaic simulator with automatic differentiation

created at May 20, 2019, 8:23 p.m.

6 +0

53 +0

16 +0

QDax by adaptive-intelligent-robotics

Accelerated Quality-Diversity

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

6 +0

243 +0

36 +0

parallax by srush

None

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

6 +0

158 +1

4 +0

PGMax by vicariousinc

Loopy belief propagation for factor graphs on discrete variables, in JAX!

created at June 9, 2021, 8:21 p.m.

6 +0

63 +0

9 +0

tf2jax by deepmind

None

created at March 2, 2022, 8:22 p.m.

7 +0

98 +2

8 +0

einshape by deepmind

None

created at July 22, 2021, 4:32 p.m.

7 +0

90 +0

5 +0

jaxdf by ucl-bug

A JAX-based research framework for writing differentiable numerical simulators with arbitrary discretizations

created at Sept. 8, 2021, 4:38 p.m.

7 +0

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

scico by lanl

Scientific Computational Imaging COde

created at Sept. 21, 2021, 11:50 a.m.

8 +0

86 +1

17 +0

jax-flows by ChrisWaites

Normalizing Flows in JAX 🌊

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

8 +0

264 +1

17 +0

SymJAX by SymJAX

Documentation:

created at Sept. 4, 2019, 2:23 p.m.

8 +0

117 -1

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

125 +1

21 +0

praxis by google

None

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

8 +0

148 +0

39 +0

extending-jax by dfm

Extending JAX with custom C++ and CUDA code

created at Jan. 7, 2021, 3:31 p.m.

9 +0

356 +1

21 +0

NuX by Information-Fusion-Lab-Umass

Normalizing Flows using JAX

created at March 9, 2020, 12:38 a.m.

9 +0

82 +0

4 +0

tinygp by dfm

The tiniest of Gaussian Process libraries

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

9 +0

282 +2

22 +0

EasyDeL by erfanzar

Accelerate your training with this open-source library. Optimize performance with streamlined training and serving options with JAX. 🚀

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

9 +0

158 -1

18 +1

jumanji by instadeepai

🕹️ A diverse suite of scalable reinforcement learning environments in JAX

created at Aug. 11, 2022, 7:34 a.m.

10 +0

539 +4

67 +0

dm_pix by deepmind

PIX is an image processing library in JAX, for JAX.

created at June 30, 2021, 4:25 p.m.

10 +0

368 +2

21 +0