jwave by ucl-bug

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

created at April 6, 2021, 10:22 a.m.

8 +0

125 +1

21 +0

jumanji by instadeepai

🕹️ A diverse suite of scalable reinforcement learning environments in JAX

created at Aug. 11, 2022, 7:34 a.m.

10 +0

539 +4

67 +0

eqxvision by paganpasta

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

created at July 24, 2022, 10:02 p.m.

4 +0

95 +0

10 +0

jaxfit by Dipolar-Quantum-Gases

GPU/TPU accelerated nonlinear least-squares curve fitting using JAX

created at Aug. 8, 2022, 5 p.m.

1 +0

43 +0

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

econpizza by gboehl

Solve nonlinear heterogeneous agent models

created at June 4, 2021, 1:34 p.m.

3 +0

63 +1

12 +0

GPJax by JaxGaussianProcesses

Gaussian processes in JAX.

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

4 +0

394 +1

47 +0

jax-md by jax-md

Differentiable, Hardware Accelerated, Molecular Dynamics

created at May 13, 2019, 9:03 p.m.

47 +0

1,106 +3

175 +0

fortuna by awslabs

A Library for Uncertainty Quantification.

created at Nov. 17, 2022, 1:11 p.m.

12 +0

855 +1

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

204 +2

91 +2

GNN-for-Combinatorial-Optimization by IvanIsCoding

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

created at Jan. 4, 2023, 11:09 p.m.

3 +0

39 +0

2 +0

jax-tqdm by jeremiecoullon

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

created at Jan. 16, 2023, 4:54 p.m.

4 +0

56 +1

6 +0

safejax by alvarobartt

Serialize JAX, Flax, Haiku, or Objax model params with 🤗`safetensors`

created at Dec. 21, 2022, 9:20 a.m.

2 +0

38 +0

2 +0

jraph by deepmind

A Graph Neural Network Library in Jax

created at Nov. 23, 2020, 10:27 a.m.

41 +0

1,329 +3

87 +0

mctx by deepmind

Monte Carlo tree search in JAX

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

29 +0

2,219 +4

171 -1

einshape by deepmind

None

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

7 +0

90 +0

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

204 +2

15 +0

optax by deepmind

Optax is a gradient processing and optimization library for JAX.

created at June 12, 2020, 3:45 p.m.

34 +0

1,510 +6

156 +2

dm-haiku by deepmind

JAX-based neural network library

created at Feb. 18, 2020, 7:14 a.m.

39 +0

2,819 +3

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

368 +2

21 +0