jax_cosmo by DifferentiableUniverseInitiative

A differentiable cosmology library in JAX

updated at May 26, 2024, 8:35 p.m.

18 +0

159 +2

31 +1

jraph by deepmind

A Graph Neural Network Library in Jax

updated at May 26, 2024, 8:34 p.m.

41 +0

1,329 +3

87 +0

EasyDeL by erfanzar

Accelerate your training with this open-source library. Optimize performance with streamlined training and serving options with JAX. 🚀

updated at May 26, 2024, 8:28 p.m.

9 +0

158 -1

18 +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 26, 2024, 8:28 p.m.

41 +0

2,257 +9

233 +0

flax by google

Flax is a neural network library for JAX that is designed for flexibility.

updated at May 26, 2024, 7:53 p.m.

84 +1

5,596 +26

595 +2

dm_pix by deepmind

PIX is an image processing library in JAX, for JAX.

updated at May 26, 2024, 7:22 p.m.

10 +0

368 +2

21 +0

mpi4jax by mpi4jax

Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python

updated at May 26, 2024, 6:38 p.m.

10 +0

379 +3

26 +0

parallax by srush

None

updated at May 26, 2024, 6:30 p.m.

6 +0

158 +1

4 +0

equinox by patrick-kidger

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

updated at May 26, 2024, 5:44 p.m.

21 +0

1,851 +7

127 -1

alphafold by deepmind

Open source code for AlphaFold.

updated at May 26, 2024, 5:37 p.m.

224 +0

11,926 +35

2,128 +7

diffrax by patrick-kidger

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

updated at May 26, 2024, 4:51 p.m.

20 +0

1,271 +11

116 +4

purejaxrl by luchris429

Really Fast End-to-End Jax RL Implementations

updated at May 26, 2024, 4:10 p.m.

12 +0

579 +1

47 +0

chex by deepmind

None

updated at May 26, 2024, 4 p.m.

17 +0

724 +2

41 +0

transformers by huggingface

🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.

updated at May 26, 2024, 2:36 p.m.

1,101 +4

126,623 +368

25,069 +70

efax by NeilGirdhar

Exponential families for JAX

updated at May 26, 2024, 12:51 p.m.

4 +0

51 +1

3 +0

dm-haiku by deepmind

JAX-based neural network library

updated at May 26, 2024, 12:50 p.m.

39 +0

2,819 +3

226 +0

tf2jax by deepmind

None

updated at May 26, 2024, 10:27 a.m.

7 +0

98 +2

8 +0

jumanji by instadeepai

🕹️ A diverse suite of scalable reinforcement learning environments in JAX

updated at May 26, 2024, 6:16 a.m.

10 +0

539 +4

67 +0

maxtext by google

A simple, performant and scalable Jax LLM!

updated at May 26, 2024, 5:06 a.m.

23 +1

1,300 +5

226 -2

SymJAX by SymJAX

Documentation:

updated at May 26, 2024, 12:52 a.m.

8 +0

117 -1

5 +0