levanter by stanford-crfm

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

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

15 +0

446 +1

64 +0

kfac-jax by deepmind

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

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

10 +0

202 +1

15 +0

tf2jax by deepmind

None

created at March 2, 2022, 8:22 p.m.

7 +0

96 +0

8 +0

mctx by deepmind

Monte Carlo tree search in JAX

created at March 1, 2022, 5:26 p.m.

29 +0

2,210 +1

172 +0

QDax by adaptive-intelligent-robotics

Accelerated Quality-Diversity

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

6 +0

243 +0

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

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

tree-math by google

Mathematical operations for JAX pytrees

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

10 -1

167 +0

7 +0

evojax by google

None

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

23 +0

789 +1

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

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

scico by lanl

Scientific Computational Imaging COde

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

8 +0

84 +0

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

lqg by RothkopfLab

Inverse optimal control for continuous psychophysics

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

2 +0

24 +0

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

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

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

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

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

einshape by deepmind

None

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

7 +0

89 +0

5 +0