praxis by google

None

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

8 +0

144 +3

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

SymJAX by SymJAX

Documentation:

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

8 +0

118 +1

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

coax by coax-dev

Modular framework for Reinforcement Learning in python

created at March 11, 2021, 12:56 a.m.

8 +0

161 +0

18 +0

equivariant-MLP by mfinzi

A library for programmatically generating equivariant layers through constraint solving

created at Sept. 11, 2020, 5:40 p.m.

9 +0

248 +1

19 +0

EasyDeL by erfanzar

EasyDeL is an OpenSource Library to make your training faster and more Optimized With cool Options for training and serving Both in Python And Mojo🔥

created at April 9, 2023, 5:18 p.m.

9 +0

157 +5

17 +0

tinygp by dfm

The tiniest of Gaussian Process libraries

created at Feb. 14, 2021, 3:51 p.m.

9 +0

280 +2

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

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

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

evosax by RobertTLange

Evolution Strategies in JAX 🦎

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

10 +0

451 +1

40 +1

tree-math by google

Mathematical operations for JAX pytrees

created at Dec. 18, 2021, 1:56 a.m.

10 -1

167 +0

7 +0

gymnax by RobertTLange

RL Environments in JAX 🌍

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

10 +0

532 +5

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

kfac-jax by deepmind

Second Order Optimization and Curvature Estimation with K-FAC in JAX.

created at March 18, 2022, 10:19 a.m.

10 +0

202 +1

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

365 +1

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

XLB by Autodesk

XLB: Accelerated Lattice Boltzmann (XLB) based on JAX for Physics-based ML

created at May 4, 2023, 5:36 p.m.

10 +0

146 +3

15 +0

kalman-jax by AaltoML

Approximate inference for Markov Gaussian processes using iterated Kalman smoothing, in JAX

created at April 3, 2020, 6:52 a.m.

11 +0

92 +1

13 +0