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

124 +1

21 +1

jumanji by instadeepai

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

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

10 +0

534 +6

67 +1

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

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

62 +3

12 +0

GPJax by JaxGaussianProcesses

Gaussian processes in JAX.

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

4 +0

389 +1

47 +0

jax-md by jax-md

Differentiable, Hardware Accelerated, Molecular Dynamics

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

46 +0

1,098 +5

175 +0

fortuna by awslabs

A Library for Uncertainty Quantification.

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

12 +0

855 +0

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

201 +0

89 +0

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

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

55 +0

4 +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,325 +2

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

einshape by deepmind

None

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

7 +0

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

202 +1

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,501 +11

152 +0

dm-haiku by deepmind

JAX-based neural network library

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

39 +0

2,816 +10

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

365 +1

21 +0