fortuna by awslabs

A Library for Uncertainty Quantification.

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

12 +0

890 -1

46 +0

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.

13 +0

850 +0

106 +0

chex by deepmind

None

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

17 +0

792 +4

44 +0

objax by google

None

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

25 +0

772 +1

77 +0

purejaxrl by luchris429

Really Fast End-to-End Jax RL Implementations

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

12 +0

740 +4

62 -1

gymnax by RobertTLange

RL Environments in JAX 🌍

created at Dec. 26, 2020, 2:35 p.m.

10 +0

659 +2

62 +0

jumanji by instadeepai

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

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

12 +0

651 +6

80 +0

netket by netket

Machine learning algorithms for many-body quantum systems

created at April 23, 2018, 6:48 p.m.

26 +0

552 +1

188 +0

BrainPy by brainpy

Brain Dynamics Programming in Python

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

12 +0

538 +2

94 +0

distrax by deepmind

None

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

18 +0

537 +0

32 +0

ott by ott-jax

Optimal transport tools implemented with the JAX framework, to get differentiable, parallel and jit-able computations.

created at Dec. 24, 2021, 7:28 a.m.

10 +0

529 +2

82 +1

levanter by stanford-crfm

Legible, Scalable, Reproducible Foundation Models with Named Tensors and Jax

created at May 24, 2022, 10:26 p.m.

14 +0

520 +1

82 +0

evosax by RobertTLange

Evolution Strategies in JAX 🦎

created at Dec. 30, 2020, 2:14 p.m.

10 +0

513 +4

43 +1

GPJax by JaxGaussianProcesses

Gaussian processes in JAX.

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

3 +0

462 +1

54 +1

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.

15 +0

458 +1

69 +1

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.

11 +0

447 +1

30 +0

vdvae by openai

Repository for the paper "Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images"

created at Nov. 10, 2020, 8:57 p.m.

128 +0

436 +0

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

390 +0

23 +0

extending-jax by dfm

Extending JAX with custom C++ and CUDA code

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

10 +0

378 +0

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

327 +2

17 +0