sklearn-jax-kernels by ExpectationMax

Composable kernels for scikit-learn implemented in JAX.

created at March 10, 2020, 6:59 p.m.

Python

5 +0

40 -1

4 +0

GitHub
jax-models by DarshanDeshpande

Unofficial JAX implementations of deep learning research papers

created at Jan. 8, 2022, 1:40 p.m.

Python

5 +0

140 +0

9 +0

GitHub
exojax by HajimeKawahara

🐈 Automatic differentiable spectrum modeling of exoplanets/brown dwarfs using JAX, compatible with NumPyro and JAXopt

created at Dec. 12, 2020, 3:46 a.m.

Jupyter Notebook

4 +0

40 +0

14 +0

GitHub
jax-tqdm by jeremiecoullon

Add a tqdm progress bar to your JAX scans and loops.

created at Jan. 16, 2023, 4:54 p.m.

Python

4 +0

55 +0

4 +0

GitHub
jaxlie by brentyi

Rigid transforms + Lie groups in JAX

created at Nov. 28, 2020, 6:18 a.m.

Python

4 +0

202 +1

14 +0

GitHub
efax by NeilGirdhar

Exponential families for JAX

created at March 24, 2020, 5:44 a.m.

Python

4 +0

50 +0

3 +0

GitHub
xmcgan_image_generation by google-research

None

created at June 24, 2021, 3:58 p.m.

Python

4 +0

98 -1

15 +0

GitHub
flaxvision by rolandgvc

A selection of neural network models ported from torchvision for JAX & Flax.

created at June 14, 2020, 4:34 p.m.

Python

4 +0

44 +0

2 +0

GitHub
eqxvision by paganpasta

A Python package of computer vision models for the Equinox ecosystem.

created at July 24, 2022, 10:02 p.m.

Python

4 +0

95 +0

10 +0

GitHub
GPJax by JaxGaussianProcesses

Gaussian processes in JAX.

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

Python

4 +0

389 +1

47 +0

GitHub
imax by 4rtemi5

Image augmentation library for Jax

created at Feb. 9, 2021, 1:49 a.m.

Python

3 +0

34 +0

3 +0

GitHub
spyx by kmheckel

Spyx: Spiking Neural Networks in JAX

created at June 21, 2023, 8:42 p.m.

Jupyter Notebook

3 +0

75 +1

9 +0

GitHub
jax-fid by matthias-wright

FID computation in Jax/Flax.

created at July 29, 2021, 4:05 a.m.

Python

3 +0

21 +0

5 +0

GitHub
econpizza by gboehl

Solve nonlinear heterogeneous agent models

created at June 4, 2021, 1:34 p.m.

Python

3 +0

62 +3

12 +0

GitHub
GNN-for-Combinatorial-Optimization by IvanIsCoding

JAX + Flax implementation of "Combinatorial Optimization with Physics-Inspired Graph Neural Networks" by Schuetz et al.

created at Jan. 4, 2023, 11:09 p.m.

Jupyter Notebook

3 +0

38 +0

2 +0

GitHub
lorax by davisyoshida

LoRA for arbitrary JAX models and functions

created at April 21, 2023, 11:50 a.m.

Python

3 +0

117 +0

4 +0

GitHub
bayex by alonfnt

Minimal Implementation of Bayesian Optimization in JAX

created at Aug. 1, 2021, 8:41 a.m.

Python

3 +0

78 +0

1 +0

GitHub
lqg by RothkopfLab

Inverse optimal control for continuous psychophysics

created at Aug. 24, 2021, 7:38 a.m.

Jupyter Notebook

2 +0

24 +0

8 +0

GitHub
JAX-Flax-Tutorial-Image-Classification-with-Linen by 8bitmp3

How to use the Flax Linen API to build a convolutional neural network model and train it for image classification (using TensorFlow Datasets).

created at Dec. 24, 2020, 3 a.m.

Jupyter Notebook

2 +0

22 +0

3 +0

GitHub
safejax by alvarobartt

Serialize JAX, Flax, Haiku, or Objax model params with 🤗`safetensors`

created at Dec. 21, 2022, 9:20 a.m.

Python

2 +0

38 +0

2 +0

GitHub