optax by deepmind

Optax is a gradient processing and optimization library for JAX.

created at June 12, 2020, 3:45 p.m.

34 +0

1,510 +6

156 +2

flaxvision by rolandgvc

A selection of neural network models ported from torchvision for JAX & Flax.

created at June 14, 2020, 4:34 p.m.

4 +0

44 +0

2 +0

BrainPy by brainpy

Brain Dynamics Programming in Python

created at July 16, 2020, 2:11 a.m.

12 +0

484 +2

88 +0

mpi4jax by mpi4jax

Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python

created at July 21, 2020, 10:57 a.m.

10 +0

379 +3

26 +0

jaxns by Joshuaalbert

Probabilistic Programming and Nested sampling in JAX

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

5 +0

129 +2

8 +0

chex by deepmind

None

created at Aug. 6, 2020, 9:32 a.m.

17 +0

724 +2

41 +0

objax by google

None

created at Aug. 20, 2020, 6:20 a.m.

26 +0

761 +0

79 +0

GPJax by JaxGaussianProcesses

Gaussian processes in JAX.

created at Sept. 27, 2020, 8:40 p.m.

4 +0

394 +1

47 +0

vdvae by openai

Repository for the paper "Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images"

created at Nov. 10, 2020, 8:57 p.m.

129 +0

429 +1

84 -1

jraph by deepmind

A Graph Neural Network Library in Jax

created at Nov. 23, 2020, 10:27 a.m.

41 +0

1,329 +3

87 +0

jaxlie by brentyi

Rigid transforms + Lie groups in JAX

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

5 +1

203 +1

14 +0

blackjax by blackjax-devs

BlackJAX is a Bayesian Inference library designed for ease of use, speed and modularity.

created at Dec. 9, 2020, 8:12 a.m.

13 +0

729 +2

94 +1

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

gymnax by RobertTLange

RL Environments in JAX 🌍

created at Dec. 26, 2020, 2:35 p.m.

10 +0

542 +7

49 +1

evosax by RobertTLange

Evolution Strategies in JAX 🦎

created at Dec. 30, 2020, 2:14 p.m.

10 +0

454 +2

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

ott by google-research

None

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

15 +0

215 +0

18 +0

imax by 4rtemi5

Image augmentation library for Jax

created at Feb. 9, 2021, 1:49 a.m.

3 +0

34 +0

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

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