tinygp by dfm

The tiniest of Gaussian Process libraries

updated at May 24, 2024, 3:46 p.m.

9 +0

282 +2

22 +0

jaxopt by google

Hardware accelerated, batchable and differentiable optimizers in JAX.

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

19 +0

895 +1

60 +0

praxis by google

None

updated at May 24, 2024, 7:32 p.m.

8 +0

148 +0

39 +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 24, 2024, 7:40 p.m.

16 +0

410 +4

61 +0

fourier-feature-networks by tancik

Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains

updated at May 24, 2024, 7:57 p.m.

24 +0

1,194 +1

128 +1

numpyro by pyro-ppl

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

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

49 +0

2,062 +5

222 +2

dynamiqs by dynamiqs

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

updated at May 25, 2024, 1:51 a.m.

6 +0

106 +1

11 +0

extending-jax by dfm

Extending JAX with custom C++ and CUDA code

updated at May 25, 2024, 4:42 a.m.

9 +0

356 +1

21 +0

levanter by stanford-crfm

Legible, Scalable, Reproducible Foundation Models with Named Tensors and Jax

updated at May 25, 2024, 8:50 a.m.

15 +0

454 +5

66 +0

mctx by deepmind

Monte Carlo tree search in JAX

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

29 +0

2,219 +4

171 -1

pytorch-fid by mseitzer

Compute FID scores with PyTorch.

updated at May 25, 2024, 12:12 p.m.

14 +0

3,146 +15

492 +1

kfac-jax by deepmind

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

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

10 +0

204 +2

15 +0

blackjax by blackjax-devs

BlackJAX is a Bayesian Inference library designed for ease of use, speed and modularity.

updated at May 25, 2024, 1:15 p.m.

13 +0

729 +2

94 +1

trax by google

Trax — Deep Learning with Clear Code and Speed

updated at May 25, 2024, 3:29 p.m.

146 +0

7,969 +6

811 +0

scenic by google-research

Scenic: A Jax Library for Computer Vision Research and Beyond

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

39 +0

3,046 +10

407 +3

optax by deepmind

Optax is a gradient processing and optimization library for JAX.

updated at May 25, 2024, 8:01 p.m.

34 +0

1,510 +6

156 +2

brax by google

Massively parallel rigidbody physics simulation on accelerator hardware.

updated at May 25, 2024, 10:01 p.m.

36 +0

2,090 +7

231 +1

cvxpylayers by cvxgrp

Differentiable convex optimization layers

updated at May 25, 2024, 10:42 p.m.

53 +0

1,730 +7

154 +0

SymJAX by SymJAX

Documentation:

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

8 +0

117 -1

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