kalman-jax by AaltoML

Approximate inference for Markov Gaussian processes using iterated Kalman smoothing, in JAX

created at April 3, 2020, 6:52 a.m.

11 +0

94 +0

13 +0

jax_cosmo by DifferentiableUniverseInitiative

A differentiable cosmology library in JAX

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

18 +0

172 +0

36 +0

big_transfer by google-research

Official repository for the "Big Transfer (BiT): General Visual Representation Learning" paper.

created at May 11, 2020, 9:34 p.m.

42 +0

1,504 +0

174 +0

parallax by srush

None

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

6 +0

156 +0

4 +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,638 +10

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

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.

25 +1

1,232 +2

131 +0

BrainPy by brainpy

Brain Dynamics Programming in Python

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

12 +0

511 +3

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

12 +0

421 +1

29 +0

jaxns by Joshuaalbert

Probabilistic Programming and Nested sampling in JAX

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

5 +0

141 +1

9 +0

chex by deepmind

None

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

18 +0

759 +1

44 +0

objax by google

None

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

26 +0

768 +0

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

252 +0

21 +0

GPJax by JaxGaussianProcesses

Gaussian processes in JAX.

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

3 +0

436 +4

51 +0

vision_transformer by google-research

None

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

102 +1

10,068 +48

1,260 +2

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.

130 +1

434 +1

84 +0

jraph by deepmind

A Graph Neural Network Library in Jax

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

40 +0

1,357 +4

88 +0

jaxlie by brentyi

Rigid transforms + Lie groups in JAX

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

6 +0

222 +0

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

797 +6

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

5 +0

45 +0

14 +0