purejaxrl by luchris429

Really Fast End-to-End Jax RL Implementations

created at Feb. 25, 2023, 3:38 p.m.

12 +0

574 +6

47 +1

GPJax by JaxGaussianProcesses

Gaussian processes in JAX.

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

4 +0

389 +1

47 +0

fortuna by awslabs

A Library for Uncertainty Quantification.

created at Nov. 17, 2022, 1:11 p.m.

12 +0

855 +0

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

evosax by RobertTLange

Evolution Strategies in JAX 🦎

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

10 +0

451 +1

40 +1

chex by deepmind

None

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

17 +0

720 +4

40 +0

praxis by google

None

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

8 +0

144 +3

39 +3

QDax by adaptive-intelligent-robotics

Accelerated Quality-Diversity

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

6 +0

243 +0

35 +0

JAX-Toolbox by NVIDIA

JAX-Toolbox

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

23 +0

177 +6

34 +1

distrax by deepmind

None

created at April 1, 2021, 5:03 p.m.

17 +0

511 +2

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

jax-unirep by ElArkk

Reimplementation of the UniRep protein featurization model.

created at Jan. 21, 2020, 3:30 p.m.

6 +0

101 +0

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

nested-transformer by google-research

Nested Hierarchical Transformer https://arxiv.org/pdf/2105.12723.pdf

created at June 1, 2021, 5:09 a.m.

8 +0

189 +0

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

tinygp by dfm

The tiniest of Gaussian Process libraries

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

9 +0

280 +2

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

365 +1

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

extending-jax by dfm

Extending JAX with custom C++ and CUDA code

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

9 +0

353 +0

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