JAX-Toolbox by NVIDIA

JAX-Toolbox

created at Feb. 13, 2023, 2:39 p.m.

23 +0

171 +4

33 +0

tree-math by google

Mathematical operations for JAX pytrees

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

11 +0

167 +0

7 +0

coax by coax-dev

Modular framework for Reinforcement Learning in python

created at March 11, 2021, 12:56 a.m.

8 +0

161 +0

18 +0

jax_cosmo by DifferentiableUniverseInitiative

A differentiable cosmology library in JAX

created at May 1, 2020, 5:47 p.m.

18 +0

157 +0

30 +0

parallax by srush

None

created at May 19, 2020, 1:49 a.m.

6 +0

157 +0

4 +0

EasyDeL by erfanzar

EasyDeL is an OpenSource Library to make your training faster and more Optimized With cool Options for training and serving Both in Python And Mojo🔥

created at April 9, 2023, 5:18 p.m.

NEW!

9 +0

152 +0

17 +0

XLB by Autodesk

XLB: Accelerated Lattice Boltzmann (XLB) based on JAX for Physics-based ML

created at May 4, 2023, 5:36 p.m.

10 +0

143 +2

15 +0

praxis by google

None

created at June 14, 2022, 4:04 p.m.

8 +0

141 +0

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

9 +0

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.

7 +0

123 +1

20 +0

jaxns by Joshuaalbert

Probabilistic Programming and Nested sampling in JAX

created at July 31, 2020, 6 p.m.

5 +0

121 +2

8 +0

lorax by davisyoshida

LoRA for arbitrary JAX models and functions

created at April 21, 2023, 11:50 a.m.

3 +0

117 +2

4 +0

SymJAX by SymJAX

Documentation:

created at Sept. 4, 2019, 2:23 p.m.

8 +0

117 +0

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

jax-unirep by ElArkk

Reimplementation of the UniRep protein featurization model.

created at Jan. 21, 2020, 3:30 p.m.

6 +0

101 +0

30 +0

xmcgan_image_generation by google-research

None

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

4 +0

99 +0

15 +0

dynamiqs by dynamiqs

High-performance quantum systems simulation with JAX (GPU-accelerated & differentiable solvers).

created at Feb. 5, 2023, 5:04 p.m.

NEW!

6 +0

98 +0

11 +0

tf2jax by deepmind

None

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

7 +0

96 +0

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

kalman-jax by AaltoML

Approximate inference for Markov Gaussian processes using iterated Kalman smoothing, in JAX

created at April 3, 2020, 6:52 a.m.

11 +0

91 +0

13 +0