exojax by HajimeKawahara

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

updated at April 27, 2024, 9:02 p.m.

4 +0

40 +0

14 +0

coax by coax-dev

Modular framework for Reinforcement Learning in python

updated at April 27, 2024, 8:39 p.m.

8 +0

161 +0

18 +0

jax_cosmo by DifferentiableUniverseInitiative

A differentiable cosmology library in JAX

updated at April 26, 2024, 4:33 p.m.

18 +0

157 +0

30 +0

tf2jax by deepmind

None

updated at April 26, 2024, 3:52 p.m.

7 +0

96 +0

8 +0

fedjax by google

FedJAX is a JAX-based open source library for Federated Learning simulations that emphasizes ease-of-use in research.

updated at April 25, 2024, 6:48 p.m.

11 +0

249 +0

41 +0

mcx by rlouf

Express & compile probabilistic programs for performant inference on CPU & GPU. Powered by JAX.

updated at April 25, 2024, 1:56 a.m.

17 +0

324 +0

17 +0

kernex by ASEM000

Stencil computations in JAX

updated at April 22, 2024, 11:47 p.m.

1 +0

59 +0

3 +0

einshape by deepmind

None

updated at April 22, 2024, 4:15 p.m.

7 +0

89 +0

5 +0

vdvae by openai

Repository for the paper "Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images"

updated at April 22, 2024, 4:14 p.m.

129 +0

427 +0

85 +0

GNN-for-Combinatorial-Optimization by IvanIsCoding

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

updated at April 22, 2024, 11:11 a.m.

3 +0

38 +0

2 +0

ott by google-research

None

updated at April 21, 2024, 10:34 p.m.

15 +0

215 +0

18 +0

QDax by adaptive-intelligent-robotics

Accelerated Quality-Diversity

updated at April 16, 2024, 8:40 a.m.

6 +0

243 +0

35 +0

jax-unirep by ElArkk

Reimplementation of the UniRep protein featurization model.

updated at April 14, 2024, 12:26 a.m.

6 +0

101 +0

30 +0

flaxvision by rolandgvc

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

updated at April 11, 2024, 2:03 p.m.

4 +0

44 +0

2 +0

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).

updated at March 31, 2024, 6:02 p.m.

2 +0

22 +0

3 +0

imax by 4rtemi5

Image augmentation library for Jax

updated at March 31, 2024, 1:22 a.m.

3 +0

34 +0

3 +0

nested-transformer by google-research

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

updated at March 28, 2024, 10:26 a.m.

8 +0

189 +0

25 +0

deltapv by romanodev

A photovoltaic simulator with automatic differentiation

updated at March 27, 2024, 8:16 a.m.

6 +0

53 +0

16 +0

NuX by Information-Fusion-Lab-Umass

Normalizing Flows using JAX

updated at March 3, 2024, 10:16 p.m.

9 +0

82 +0

4 +0

jax-fid by matthias-wright

FID computation in Jax/Flax.

updated at Jan. 4, 2024, 4:59 p.m.

3 +0

21 +0

5 +0