gymnax by RobertTLange

RL Environments in JAX 🌍

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

10 +0

532 +5

48 +0

JAX-Flax-Tutorial-Image-Classification-with-Linen by 8bitmp3

How to use the Flax Linen API to build a convolutional neural network model and train it for image classification (using TensorFlow Datasets).

created at Dec. 24, 2020, 3 a.m.

2 +0

22 +0

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

cr-sparse by carnotresearch

Functional models and algorithms for sparse signal processing

created at Dec. 22, 2020, 8:31 a.m.

6 +0

83 +0

11 +0

exojax by HajimeKawahara

🐈 Automatic differentiable spectrum modeling of exoplanets/brown dwarfs using JAX, compatible with NumPyro and JAXopt

created at Dec. 12, 2020, 3:46 a.m.

4 +0

40 +0

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.

14 +0

727 +0

93 +3

jaxlie by brentyi

Rigid transforms + Lie groups in JAX

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

4 +0

202 +1

14 +0

jraph by deepmind

A Graph Neural Network Library in Jax

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

41 +0

1,325 +2

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

427 +0

85 +0

vision_transformer by google-research

None

created at Oct. 21, 2020, 12:35 p.m.

95 -1

9,367 +45

1,206 +3

GPJax by JaxGaussianProcesses

Gaussian processes in JAX.

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

4 +0

389 +1

47 +0

equivariant-MLP by mfinzi

A library for programmatically generating equivariant layers through constraint solving

created at Sept. 11, 2020, 5:40 p.m.

9 +0

248 +1

19 +0

objax by google

None

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

26 +0

761 +0

79 +0

chex by deepmind

None

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

17 +0

720 +4

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

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.

9 +0

373 +1

26 +1

BrainPy by brainpy

Brain Dynamics Programming in Python

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

12 +0

480 +1

88 +1

fourier-feature-networks by tancik

Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains

created at June 16, 2020, 11:54 p.m.

24 +0

1,190 +2

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

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,501 +11

152 +0