levanter by stanford-crfm

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

created at May 24, 2022, 10:26 p.m.

14 +0

520 +1

82 +0

evosax by RobertTLange

Evolution Strategies in JAX 🦎

created at Dec. 30, 2020, 2:14 p.m.

10 +0

513 +4

43 +1

GPJax by JaxGaussianProcesses

Gaussian processes in JAX.

created at Sept. 27, 2020, 8:40 p.m.

3 +0

462 +1

54 +1

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.

created at June 14, 2022, 4:04 p.m.

15 +0

458 +1

69 +1

mpi4jax by mpi4jax

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

created at July 21, 2020, 10:57 a.m.

11 +0

447 +1

30 +0

vdvae by openai

Repository for the paper "Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images"

created at Nov. 10, 2020, 8:57 p.m.

128 +0

436 +0

86 +0

dm_pix by deepmind

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

created at June 30, 2021, 4:25 p.m.

10 +0

390 +0

23 +0

extending-jax by dfm

Extending JAX with custom C++ and CUDA code

created at Jan. 7, 2021, 3:31 p.m.

10 +0

378 +0

23 +0

mcx by rlouf

Express & compile probabilistic programs for performant inference on CPU & GPU. Powered by JAX.

created at Jan. 22, 2020, 8:38 a.m.

17 +0

327 +2

17 +0

tinygp by dfm

The tiniest of Gaussian Process libraries

created at Feb. 14, 2021, 3:51 p.m.

10 +1

298 +1

24 +0

jax-flows by ChrisWaites

Normalizing Flows in JAX 🌊

created at March 23, 2020, 2:42 a.m.

8 +0

274 +0

19 +0

QDax by adaptive-intelligent-robotics

Accelerated Quality-Diversity

created at Feb. 11, 2022, 3:48 p.m.

5 +0

268 +0

46 +0

equivariant-MLP by mfinzi

A library for programmatically generating equivariant layers through constraint solving

created at Sept. 11, 2020, 5:40 p.m.

10 +0

257 +0

22 +1

fedjax by google

FedJAX is a JAX-based open source library for Federated Learning simulations that emphasizes ease-of-use in research.

created at Dec. 22, 2020, 2:44 p.m.

11 +0

254 +0

41 +0

JAX-Toolbox by NVIDIA

JAX-Toolbox

created at Feb. 13, 2023, 2:39 p.m.

22 +0

254 +4

48 -1

kfac-jax by deepmind

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

created at March 18, 2022, 10:19 a.m.

11 +0

251 +1

23 +0

spu by secretflow

SPU (Secure Processing Unit) aims to be a provable, measurable secure computation device, which provides computation ability while keeping your private data protected.

created at July 2, 2022, 9:05 a.m.

8 +0

243 +0

106 +1

XLB by Autodesk

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

created at May 4, 2023, 5:36 p.m.

12 +0

238 +2

25 +0

flaxmodels by matthias-wright

Pretrained deep learning models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet, etc.

created at May 16, 2021, 1:34 p.m.

6 +0

238 +0

24 +0

jaxlie by brentyi

Rigid transforms + Lie groups in JAX

created at Nov. 28, 2020, 6:18 a.m.

7 +1

234 +0

15 +0