xmcgan_image_generation by google-research

None

created at June 24, 2021, 3:58 p.m.

4 +0

98 -1

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

365 +1

21 +0

scenic by google-research

Scenic: A Jax Library for Computer Vision Research and Beyond

created at July 12, 2021, 2:27 p.m.

39 +0

3,029 +12

401 +0

jaxopt by google

Hardware accelerated, batchable and differentiable optimizers in JAX.

created at July 12, 2021, 5:16 p.m.

19 +0

891 +2

59 +1

einshape by deepmind

None

created at July 22, 2021, 4:32 p.m.

7 +0

89 +0

5 +0

diffrax by patrick-kidger

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

created at July 25, 2021, 11:42 p.m.

20 +0

1,251 +9

111 +1

equinox by patrick-kidger

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

created at July 29, 2021, 2:21 a.m.

21 +0

1,834 +9

127 +0

jax-fid by matthias-wright

FID computation in Jax/Flax.

created at July 29, 2021, 4:05 a.m.

3 +0

21 +0

5 +0

bayex by alonfnt

Minimal Implementation of Bayesian Optimization in JAX

created at Aug. 1, 2021, 8:41 a.m.

3 +0

78 +0

1 +0

two-player-auctions by degregat

JAX/Haiku implementation of "Auction Learning as a Two-Player Game"

created at Aug. 23, 2021, 8:04 a.m.

0 +0

0 +0

0 +0

lqg by RothkopfLab

Inverse optimal control for continuous psychophysics

created at Aug. 24, 2021, 7:38 a.m.

2 +0

24 +0

8 +0

jaxdf by ucl-bug

A JAX-based research framework for writing differentiable numerical simulators with arbitrary discretizations

created at Sept. 8, 2021, 4:38 p.m.

7 +0

107 +0

7 +0

scico by lanl

Scientific Computational Imaging COde

created at Sept. 21, 2021, 11:50 a.m.

8 +0

84 +0

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

created at Oct. 28, 2021, 8 p.m.

8 +0

569 +4

92 +0

gumbel_max_causal_gadgets_part2 by GuyLor

Part 2 of the code for the NeurIPS 2021 paper "Learning Generalized Gumbel-max Causal Mechanisms"

created at Nov. 12, 2021, 6:16 p.m.

1 +0

2 +0

0 +0

evojax by google

None

created at Dec. 7, 2021, 12:30 a.m.

23 +0

789 +1

69 +0

tree-math by google

Mathematical operations for JAX pytrees

created at Dec. 18, 2021, 1:56 a.m.

10 -1

167 +0

7 +0

ott by ott-jax

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

created at Dec. 24, 2021, 7:28 a.m.

10 +0

456 +0

79 +0

jax-models by DarshanDeshpande

Unofficial JAX implementations of deep learning research papers

created at Jan. 8, 2022, 1:40 p.m.

5 +0

140 +0

9 +0

QDax by adaptive-intelligent-robotics

Accelerated Quality-Diversity

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

6 +0

243 +0

35 +0