GPJax by JaxGaussianProcesses

Gaussian processes in JAX.

updated at May 21, 2024, 6:06 a.m.

4 +0

394 +1

47 +0

vdvae by openai

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

updated at May 21, 2024, 2:14 a.m.

129 +0

429 +1

84 -1

econpizza by gboehl

Solve nonlinear heterogeneous agent models

updated at May 21, 2024, 1:47 a.m.

3 +0

63 +1

12 +0

jax-tqdm by jeremiecoullon

Add a tqdm progress bar to your JAX scans and loops.

updated at May 21, 2024, 12:05 a.m.

4 +0

56 +1

6 +0

scico by lanl

Scientific Computational Imaging COde

updated at May 20, 2024, 10:59 p.m.

8 +0

86 +1

17 +0

jwave by ucl-bug

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

updated at May 20, 2024, 12:52 p.m.

8 +0

125 +1

21 +0

XLB by Autodesk

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

updated at May 20, 2024, 12:52 p.m.

10 +0

147 +1

14 -1

kalman-jax by AaltoML

Approximate inference for Markov Gaussian processes using iterated Kalman smoothing, in JAX

updated at May 18, 2024, 2:42 p.m.

11 +0

93 +0

13 +0

lorax by davisyoshida

LoRA for arbitrary JAX models and functions

updated at May 18, 2024, 10:13 a.m.

3 +0

119 +0

4 +0

mipnerf by google

None

updated at May 18, 2024, 6:03 a.m.

18 +0

885 +0

108 +1

ott by ott-jax

Optimal transport tools implemented with the JAX framework, to get differentiable, parallel and jit-able computations.

updated at May 17, 2024, 6:39 p.m.

10 +0

458 +0

78 +0

tensorly by tensorly

TensorLy: Tensor Learning in Python.

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

45 +0

1,504 +0

281 +0

tree-math by google

Mathematical operations for JAX pytrees

updated at May 16, 2024, 1:55 p.m.

10 +0

169 +0

7 +0

fedjax by google

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

updated at May 15, 2024, 7:39 p.m.

11 +0

249 +0

41 +0

einshape by deepmind

None

updated at May 13, 2024, 7:39 p.m.

7 +0

90 +0

5 +0

distrax by deepmind

None

updated at May 13, 2024, 5:01 p.m.

17 +0

512 +0

32 +0

GNN-for-Combinatorial-Optimization by IvanIsCoding

JAX + Flax implementation of "Combinatorial Optimization with Physics-Inspired Graph Neural Networks" by Schuetz et al.

updated at May 13, 2024, 1:51 p.m.

3 +0

39 +0

2 +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 +0

4 +0

flaxmodels by matthias-wright

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

updated at May 9, 2024, 12:33 p.m.

6 +0

223 +0

24 +0

eqxvision by paganpasta

A Python package of computer vision models for the Equinox ecosystem.

updated at May 8, 2024, 8:14 p.m.

4 +0

95 +0

10 +0