jax-tqdm by jeremiecoullon

Add a tqdm progress bar to your JAX scans and loops.

created at Jan. 16, 2023, 4:54 p.m.

4 +0

55 +0

4 +0

jaxlie by brentyi

Rigid transforms + Lie groups in JAX

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

4 +0

202 +1

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

jax-models by DarshanDeshpande

Unofficial JAX implementations of deep learning research papers

created at Jan. 8, 2022, 1:40 p.m.

5 +0

140 +0

9 +0

sklearn-jax-kernels by ExpectationMax

Composable kernels for scikit-learn implemented in JAX.

created at March 10, 2020, 6:59 p.m.

5 +0

40 -1

4 +0

deltapv by romanodev

A photovoltaic simulator with automatic differentiation

created at May 20, 2019, 8:23 p.m.

6 +0

53 +0

16 +0

parallax by srush

None

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

6 +0

157 +0

4 +0

dynamiqs by dynamiqs

High-performance quantum systems simulation with JAX (GPU-accelerated & differentiable solvers).

created at Feb. 5, 2023, 5:04 p.m.

6 +0

98 +0

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

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

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

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

QDax by adaptive-intelligent-robotics

Accelerated Quality-Diversity

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

6 +0

243 +0

35 +0

einshape by deepmind

None

created at July 22, 2021, 4:32 p.m.

7 +0

89 +0

5 +0

tf2jax by deepmind

None

created at March 2, 2022, 8:22 p.m.

7 +0

96 +0

8 +0

jaxdf by ucl-bug

A JAX-based research framework for writing differentiable numerical simulators with arbitrary discretizations

created at Sept. 8, 2021, 4:38 p.m.

7 +0

107 +0

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

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

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

jax-flows by ChrisWaites

Normalizing Flows in JAX 🌊

created at March 23, 2020, 2:42 a.m.

8 +0

263 +0

17 +0