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 10, 2024, 10:47 a.m.

42 +0

2,247 +1

233 +0

tensorly by tensorly

TensorLy: Tensor Learning in Python.

updated at May 10, 2024, 12:14 p.m.

45 +0

1,501 +4

281 +0

scenic by google-research

Scenic: A Jax Library for Computer Vision Research and Beyond

updated at May 10, 2024, 1:34 p.m.

39 +0

3,029 +12

401 +0

mpi4jax by mpi4jax

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

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

9 +0

373 +1

26 +1

sklearn-jax-kernels by ExpectationMax

Composable kernels for scikit-learn implemented in JAX.

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

5 +0

40 -1

4 +0

GPJax by JaxGaussianProcesses

Gaussian processes in JAX.

updated at May 11, 2024, 12:08 a.m.

4 +0

389 +1

47 +0

maxtext by google

A simple, performant and scalable Jax LLM!

updated at May 11, 2024, 1:29 a.m.

22 +0

1,281 +9

228 +1

dynamiqs by dynamiqs

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

updated at May 11, 2024, 1:55 a.m.

6 +0

98 +0

11 +0

spyx by kmheckel

Spyx: Spiking Neural Networks in JAX

updated at May 11, 2024, 2:54 a.m.

3 +0

75 +1

9 +0

XLB by Autodesk

XLB: Accelerated Lattice Boltzmann (XLB) based on JAX for Physics-based ML

updated at May 11, 2024, 4:15 a.m.

10 +0

146 +3

15 +0

praxis by google

None

updated at May 11, 2024, 7:07 a.m.

8 +0

144 +3

39 +3

jwave by ucl-bug

A JAX-based research framework for differentiable and parallelizable acoustic simulations, on CPU, GPUs and TPUs

updated at May 11, 2024, 7:19 a.m.

8 +1

124 +1

21 +1

jraph by deepmind

A Graph Neural Network Library in Jax

updated at May 11, 2024, 7:48 a.m.

41 +0

1,325 +2

85 +0

brax by google

Massively parallel rigidbody physics simulation on accelerator hardware.

updated at May 11, 2024, 7:49 a.m.

36 +0

2,077 +4

230 +0

econpizza by gboehl

Solve nonlinear heterogeneous agent models

updated at May 11, 2024, 9:01 a.m.

3 +0

62 +3

12 +0

rlax by deepmind

None

updated at May 11, 2024, 9:17 a.m.

34 +0

1,190 +4

84 +1

SymJAX by SymJAX

Documentation:

updated at May 11, 2024, 10:16 a.m.

8 +0

118 +1

5 +0

paxml by google

Pax is a Jax-based machine learning framework for training large scale models. Pax allows for advanced and fully configurable experimentation and parallelization, and has demonstrated industry leading model flop utilization rates.

updated at May 11, 2024, 10:33 a.m.

16 +1

399 +7

60 +3

tinygp by dfm

The tiniest of Gaussian Process libraries

updated at May 11, 2024, 12:01 p.m.

9 +0

280 +2

22 +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 11, 2024, 12:45 p.m.

8 +0

569 +4

92 +0