netket by netket

Machine learning algorithms for many-body quantum systems

updated at May 5, 2024, 8:13 p.m.

25 +0

502 +1

171 +0

mctx by deepmind

Monte Carlo tree search in JAX

updated at May 5, 2024, 8:02 p.m.

29 +0

2,209 +6

172 +0

tensorly by tensorly

TensorLy: Tensor Learning in Python.

updated at May 5, 2024, 7:16 p.m.

45 +0

1,497 +3

281 +0

vision_transformer by google-research

None

updated at May 5, 2024, 7:11 p.m.

96 +0

9,322 +33

1,203 +3

scenic by google-research

Scenic: A Jax Library for Computer Vision Research and Beyond

updated at May 5, 2024, 6:17 p.m.

39 +0

3,017 +18

401 +4

equinox by patrick-kidger

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

updated at May 5, 2024, 6 p.m.

21 -1

1,825 +13

127 +3

diffrax by patrick-kidger

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

updated at May 5, 2024, 4:25 p.m.

20 +0

1,242 +8

110 +0

alphafold by deepmind

Open source code for AlphaFold.

updated at May 5, 2024, 3:21 p.m.

221 +2

11,707 +23

2,081 +1

pytorch-fid by mseitzer

Compute FID scores with PyTorch.

updated at May 5, 2024, 3:15 p.m.

14 +0

3,085 +13

492 +1

mipnerf by google

None

updated at May 5, 2024, 2:53 p.m.

18 +0

881 +2

107 +0

get-started-with-JAX by gordicaleksa

The purpose of this repo is to make it easy to get started with JAX, Flax, and Haiku. It contains my "Machine Learning with JAX" series of tutorials (YouTube videos and Jupyter Notebooks) as well as the content I found useful while learning about the JAX ecosystem.

updated at May 5, 2024, 2:09 p.m.

8 +0

565 +7

92 +1

purejaxrl by luchris429

Really Fast End-to-End Jax RL Implementations

updated at May 5, 2024, 1:54 p.m.

12 +0

568 +9

46 +0

cr-sparse by carnotresearch

Functional models and algorithms for sparse signal processing

updated at May 5, 2024, 1:49 p.m.

6 +0

83 +3

11 +1

dynamiqs by dynamiqs

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

updated at May 5, 2024, 11:39 a.m.

NEW!

6 +0

98 +0

11 +0

kfac-jax by deepmind

Second Order Optimization and Curvature Estimation with K-FAC in JAX.

updated at May 5, 2024, 11:23 a.m.

10 +1

201 +4

15 +0

EasyDeL by erfanzar

EasyDeL is an OpenSource Library to make your training faster and more Optimized With cool Options for training and serving Both in Python And Mojo🔥

updated at May 5, 2024, 10:34 a.m.

NEW!

9 +0

152 +0

17 +0

numpyro by pyro-ppl

Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.

updated at May 5, 2024, 9:14 a.m.

49 +0

2,048 +8

220 +1

EasyLM by young-geng

Large language models (LLMs) made easy, EasyLM is a one stop solution for pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.

updated at May 5, 2024, 5:31 a.m.

42 +1

2,246 +9

233 +1

jax-flows by ChrisWaites

Normalizing Flows in JAX 🌊

updated at May 5, 2024, 4:30 a.m.

8 +0

263 +1

17 +0

jaxlie by brentyi

Rigid transforms + Lie groups in JAX

updated at May 5, 2024, 3:42 a.m.

4 +0

201 +1

14 +0