optax by deepmind

Optax is a gradient processing and optimization library for JAX.

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

34 +0

1,501 +11

152 +0

mipnerf by google

None

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

18 +0

883 +2

107 +0

diffrax by patrick-kidger

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

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

20 +0

1,251 +9

111 +1

mctx by deepmind

Monte Carlo tree search in JAX

updated at May 11, 2024, 1:58 p.m.

29 +0

2,210 +1

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

tinygp by dfm

The tiniest of Gaussian Process libraries

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

9 +0

280 +2

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

SymJAX by SymJAX

Documentation:

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

8 +0

118 +1

5 +0

rlax by deepmind

None

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

34 +0

1,190 +4

84 +1

econpizza by gboehl

Solve nonlinear heterogeneous agent models

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

3 +0

62 +3

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

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

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

praxis by google

None

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

8 +0

144 +3

39 +3

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

spyx by kmheckel

Spyx: Spiking Neural Networks in JAX

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

3 +0

75 +1

9 +0

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

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

GPJax by JaxGaussianProcesses

Gaussian processes in JAX.

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

4 +0

389 +1

47 +0

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