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

econpizza by gboehl

Solve nonlinear heterogeneous agent models

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

3 +0

63 +1

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

GPJax by JaxGaussianProcesses

Gaussian processes in JAX.

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

4 +0

394 +1

47 +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 21, 2024, 4:43 p.m.

8 +0

574 +2

92 +0

equivariant-MLP by mfinzi

A library for programmatically generating equivariant layers through constraint solving

updated at May 21, 2024, 6:27 p.m.

9 +0

249 +1

19 +0

cr-sparse by carnotresearch

Functional models and algorithms for sparse signal processing

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

6 +0

85 +1

11 +0

kernex by ASEM000

Stencil computations in JAX

updated at May 22, 2024, 6:04 p.m.

1 +0

62 +1

3 +0

jaxns by Joshuaalbert

Probabilistic Programming and Nested sampling in JAX

updated at May 23, 2024, 8:29 a.m.

5 +0

129 +2

8 +0

jaxlie by brentyi

Rigid transforms + Lie groups in JAX

updated at May 23, 2024, 12:38 p.m.

5 +1

203 +1

14 +0

fortuna by awslabs

A Library for Uncertainty Quantification.

updated at May 23, 2024, 1:46 p.m.

12 +0

855 +1

45 +0

jax-flows by ChrisWaites

Normalizing Flows in JAX 🌊

updated at May 23, 2024, 7:21 p.m.

8 +0

264 +1

17 +0

JAX-Toolbox by NVIDIA

JAX-Toolbox

updated at May 24, 2024, 8:51 a.m.

23 +0

184 +2

34 +0

evosax by RobertTLange

Evolution Strategies in JAX 🦎

updated at May 24, 2024, 9:36 a.m.

10 +0

454 +2

40 +0

BrainPy by brainpy

Brain Dynamics Programming in Python

updated at May 24, 2024, 10:53 a.m.

12 +0

484 +2

88 +0

big_transfer by google-research

Official repository for the "Big Transfer (BiT): General Visual Representation Learning" paper.

updated at May 24, 2024, 10:54 a.m.

42 +0

1,495 +0

174 +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.

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

8 +0

204 +2

91 +2

rlax by deepmind

None

updated at May 24, 2024, 12:07 p.m.

34 +0

1,194 +2

83 +0

gymnax by RobertTLange

RL Environments in JAX 🌍

updated at May 24, 2024, 12:37 p.m.

10 +0

542 +7

49 +1

netket by netket

Machine learning algorithms for many-body quantum systems

updated at May 24, 2024, 12:42 p.m.

25 +0

505 +2

174 -1