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

110 +0

fourier-feature-networks by tancik

Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains

created at June 16, 2020, 11:54 p.m.

24 +0

1,188 +1

125 +0

rlax by deepmind

None

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

34 +0

1,186 +3

83 +0

jax-md by jax-md

Differentiable, Hardware Accelerated, Molecular Dynamics

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

46 +0

1,093 +1

175 +2

jaxopt by google

Hardware accelerated, batchable and differentiable optimizers in JAX.

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

19 +0

889 +0

58 +0

mipnerf by google

None

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

18 +0

881 +2

107 +0

fortuna by awslabs

A Library for Uncertainty Quantification.

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

12 +0

855 +4

45 +0

evojax by google

None

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

23 +0

788 +2

69 +0

objax by google

None

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

26 +0

761 +1

79 +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.

14 +0

727 +1

90 +0

chex by deepmind

None

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

17 +0

716 +2

40 +0

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

579 +2

62 +0

purejaxrl by luchris429

Really Fast End-to-End Jax RL Implementations

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

12 +0

568 +9

46 +0

get-started-with-JAX by gordicaleksa

The purpose of this repo is to make it easy to get started with JAX, Flax, and Haiku. It contains my "Machine Learning with JAX" series of tutorials (YouTube videos and Jupyter Notebooks) as well as the content I found useful while learning about the JAX ecosystem.

created at Oct. 28, 2021, 8 p.m.

8 +0

565 +7

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

528 +4

66 +0

gymnax by RobertTLange

RL Environments in JAX 🌍

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

10 +0

527 +4

48 +1

distrax by deepmind

None

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

17 +0

509 +1

32 +0

netket by netket

Machine learning algorithms for many-body quantum systems

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

25 +0

502 +1

171 +0

BrainPy by brainpy

Brain Dynamics Programming in Python

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

12 +0

479 +4

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

456 +0

79 +13