trax by google

Trax — Deep Learning with Clear Code and Speed

created at Oct. 5, 2019, 3:09 p.m.

146 +0

7,963 +4

810 +0

big_transfer by google-research

Official repository for the "Big Transfer (BiT): General Visual Representation Learning" paper.

created at May 11, 2020, 9:34 p.m.

42 +0

1,493 +0

174 +0

efax by NeilGirdhar

Exponential families for JAX

created at March 24, 2020, 5:44 a.m.

4 +0

50 +0

3 +0

NuX by Information-Fusion-Lab-Umass

Normalizing Flows using JAX

created at March 9, 2020, 12:38 a.m.

9 +0

82 +0

4 +0

flax by google

Flax is a neural network library for JAX that is designed for flexibility.

created at Jan. 10, 2020, 9:48 a.m.

84 +0

5,546 +8

594 +1

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

fedjax by google

FedJAX is a JAX-based open source library for Federated Learning simulations that emphasizes ease-of-use in research.

created at Dec. 22, 2020, 2:44 p.m.

11 +0

249 +0

41 +0

mcx by rlouf

Express & compile probabilistic programs for performant inference on CPU & GPU. Powered by JAX.

created at Jan. 22, 2020, 8:38 a.m.

17 +0

324 +0

17 +0

jax-flows by ChrisWaites

Normalizing Flows in JAX 🌊

created at March 23, 2020, 2:42 a.m.

8 +0

263 +0

17 +0

flaxvision by rolandgvc

A selection of neural network models ported from torchvision for JAX & Flax.

created at June 14, 2020, 4:34 p.m.

4 +0

44 +0

2 +0

extending-jax by dfm

Extending JAX with custom C++ and CUDA code

created at Jan. 7, 2021, 3:31 p.m.

9 +0

353 +0

21 +0

sklearn-jax-kernels by ExpectationMax

Composable kernels for scikit-learn implemented in JAX.

created at March 10, 2020, 6:59 p.m.

5 +0

40 -1

4 +0

objax by google

None

created at Aug. 20, 2020, 6:20 a.m.

26 +0

761 +0

79 +0

imax by 4rtemi5

Image augmentation library for Jax

created at Feb. 9, 2021, 1:49 a.m.

3 +0

34 +0

3 +0

numpyro by pyro-ppl

Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.

created at Feb. 13, 2019, 9:13 p.m.

49 +0

2,052 +4

221 +1

jaxns by Joshuaalbert

Probabilistic Programming and Nested sampling in JAX

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

5 +0

125 +4

8 +0

SymJAX by SymJAX

Documentation:

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

8 +0

118 +1

5 +0

ott by google-research

None

created at Jan. 18, 2021, 3:03 p.m.

15 +0

215 +0

18 +0

mpi4jax by mpi4jax

Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python

created at July 21, 2020, 10:57 a.m.

9 +0

373 +1

26 +1

transformers by huggingface

🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.

created at Oct. 29, 2018, 1:56 p.m.

1,094 +1

125,856 +423

24,937 +74