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

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,190 +2

127 +2

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

mipnerf by google

None

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

18 +0

883 +2

107 +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 +0

93 +3

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

569 +4

92 +0

spu by secretflow

SPU (Secure Processing Unit) aims to be a provable, measurable secure computation device, which provides computation ability while keeping your private data protected.

created at July 2, 2022, 9:05 a.m.

8 +0

201 +0

89 +0

BrainPy by brainpy

Brain Dynamics Programming in Python

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

12 +0

480 +1

88 +1

jraph by deepmind

A Graph Neural Network Library in Jax

created at Nov. 23, 2020, 10:27 a.m.

41 +0

1,325 +2

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

129 +0

427 +0

85 +0

rlax by deepmind

None

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

34 +0

1,190 +4

84 +1

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

objax by google

None

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

26 +0

761 +0

79 +0

evojax by google

None

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

23 +0

789 +1

69 +0

jumanji by instadeepai

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

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

10 +0

534 +6

67 +1

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

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

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

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

gymnax by RobertTLange

RL Environments in JAX 🌍

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

10 +0

532 +5

48 +0