nested-transformer by google-research

Nested Hierarchical Transformer https://arxiv.org/pdf/2105.12723.pdf

created at June 1, 2021, 5:09 a.m.

8 +0

189 +0

25 +0

jwave by ucl-bug

A JAX-based research framework for differentiable and parallelizable acoustic simulations, on CPU, GPUs and TPUs

created at April 6, 2021, 10:22 a.m.

8 +1

124 +1

21 +1

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

SymJAX by SymJAX

Documentation:

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

8 +0

118 +1

5 +0

praxis by google

None

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

8 +0

144 +3

39 +3

scico by lanl

Scientific Computational Imaging COde

created at Sept. 21, 2021, 11:50 a.m.

8 +0

84 +0

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

569 +4

92 +0

jaxdf by ucl-bug

A JAX-based research framework for writing differentiable numerical simulators with arbitrary discretizations

created at Sept. 8, 2021, 4:38 p.m.

7 +0

107 +0

7 +0

einshape by deepmind

None

created at July 22, 2021, 4:32 p.m.

7 +0

89 +0

5 +0

tf2jax by deepmind

None

created at March 2, 2022, 8:22 p.m.

7 +0

96 +0

8 +0

PGMax by vicariousinc

Loopy belief propagation for factor graphs on discrete variables, in JAX!

created at June 9, 2021, 8:21 p.m.

6 +0

63 +0

9 +0

jax-unirep by ElArkk

Reimplementation of the UniRep protein featurization model.

created at Jan. 21, 2020, 3:30 p.m.

6 +0

101 +0

30 +0

cr-sparse by carnotresearch

Functional models and algorithms for sparse signal processing

created at Dec. 22, 2020, 8:31 a.m.

6 +0

83 +0

11 +0

parallax by srush

None

created at May 19, 2020, 1:49 a.m.

6 +0

157 +0

4 +0

QDax by adaptive-intelligent-robotics

Accelerated Quality-Diversity

created at Feb. 11, 2022, 3:48 p.m.

6 +0

243 +0

35 +0

dynamiqs by dynamiqs

High-performance quantum systems simulation with JAX (GPU-accelerated & differentiable solvers).

created at Feb. 5, 2023, 5:04 p.m.

6 +0

98 +0

11 +0

deltapv by romanodev

A photovoltaic simulator with automatic differentiation

created at May 20, 2019, 8:23 p.m.

6 +0

53 +0

16 +0

flaxmodels by matthias-wright

Pretrained deep learning models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet, etc.

created at May 16, 2021, 1:34 p.m.

6 +0

223 +1

24 +0

jaxns by Joshuaalbert

Probabilistic Programming and Nested sampling in JAX

created at July 31, 2020, 6 p.m.

5 +0

125 +4

8 +0