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

fortuna by awslabs

A Library for Uncertainty Quantification.

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

12 +0

855 +0

45 +0

purejaxrl by luchris429

Really Fast End-to-End Jax RL Implementations

created at Feb. 25, 2023, 3:38 p.m.

12 +0

574 +6

47 +1

jaxrl by ikostrikov

JAX (Flax) implementation of algorithms for Deep Reinforcement Learning with continuous action spaces.

created at Jan. 16, 2021, 1:26 a.m.

12 +0

583 +4

62 +0

BrainPy by brainpy

Brain Dynamics Programming in Python

created at July 16, 2020, 2:11 a.m.

12 +0

480 +1

88 +1

blackjax by blackjax-devs

BlackJAX is a Bayesian Inference library designed for ease of use, speed and modularity.

created at Dec. 9, 2020, 8:12 a.m.

14 +0

727 +0

93 +3

pytorch-fid by mseitzer

Compute FID scores with PyTorch.

created at Feb. 10, 2018, 12:33 p.m.

14 +0

3,116 +31

490 -2

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

ott by google-research

None

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

15 +0

215 +0

18 +0

paxml by google

Pax is a Jax-based machine learning framework for training large scale models. Pax allows for advanced and fully configurable experimentation and parallelization, and has demonstrated industry leading model flop utilization rates.

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

16 +1

399 +7

60 +3

distrax by deepmind

None

created at April 1, 2021, 5:03 p.m.

17 +0

511 +2

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

chex by deepmind

None

created at Aug. 6, 2020, 9:32 a.m.

17 +0

720 +4

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

mipnerf by google

None

created at May 14, 2021, 4:01 p.m.

18 +0

883 +2

107 +0

jaxopt by google

Hardware accelerated, batchable and differentiable optimizers in JAX.

created at July 12, 2021, 5:16 p.m.

19 +0

891 +2

59 +1

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

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

maxtext by google

A simple, performant and scalable Jax LLM!

created at Feb. 28, 2023, 7:47 p.m.

22 +0

1,281 +9

228 +1

evojax by google

None

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

23 +0

789 +1

69 +0