jaxopt by google

Hardware accelerated, batchable and differentiable optimizers in JAX.

created at July 12, 2021, 5:16 p.m.

19 +0

895 +1

60 +0

scenic by google-research

Scenic: A Jax Library for Computer Vision Research and Beyond

created at July 12, 2021, 2:27 p.m.

39 +0

3,046 +10

407 +3

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

xmcgan_image_generation by google-research

None

created at June 24, 2021, 3:58 p.m.

4 +0

98 +0

15 +0

alphafold by deepmind

Open source code for AlphaFold.

created at June 17, 2021, 2:06 p.m.

224 +0

11,926 +35

2,128 +7

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

econpizza by gboehl

Solve nonlinear heterogeneous agent models

created at June 4, 2021, 1:34 p.m.

3 +0

63 +1

12 +0

brax by google

Massively parallel rigidbody physics simulation on accelerator hardware.

created at June 2, 2021, 9:45 p.m.

36 +0

2,090 +7

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

24 +0

mipnerf by google

None

created at May 14, 2021, 4:01 p.m.

18 +0

885 +0

108 +1

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

distrax by deepmind

None

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

17 +0

512 +0

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

tinygp by dfm

The tiniest of Gaussian Process libraries

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

9 +0

282 +2

22 +0

imax by 4rtemi5

Image augmentation library for Jax

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

3 +0

34 +0

3 +0

ott by google-research

None

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

15 +0

215 +0

18 +0

jaxrl by ikostrikov

JAX (Flax) implementation of algorithms for Deep Reinforcement Learning with continuous action spaces.

created at Jan. 16, 2021, 1:26 a.m.

12 +0

588 +4

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

evosax by RobertTLange

Evolution Strategies in JAX 🦎

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

10 +0

454 +2

40 +0