JAX
216 items from n2cholas/awesome-jax ★2,166
-
-
Trax github.com
"Batteries included" deep learning library focused on providing solutions for common workloads.
-
-
AlphaFold github.com
Implementation of the inference pipeline of AlphaFold v2.0, presented in Highly accurate protein structure prediction with AlphaFold.
-
Vision Transformer github.com
Official implementation of An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.
-
Big Transfer (BiT) github.com
Implementation of Big Transfer (BiT): General Visual Representation Learning.
-
-
-
BRAX github.com
Differentiable physics engine to simulate environments along with learning algorithms to train agents for these environments.
-
Equinox github.com
Callable PyTrees and filtered JIT/grad transformations => neural networks in JAX.
-
-
-
EasyLM github.com
LLMs made easy: Pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.
-
MaxText github.com
A simple, performant and scalable Jax LLM written in pure Python/Jax and targeting Google Cloud TPUs.
-
Neural Tangents github.com
High-level API for specifying neural networks of both finite and infinite width.
-
-
-
-
Penzai github.com
Prioritizes legibility, visualization, and easy editing of neural network models with composable tools and a simple mental model.
-
-
-
-
-
Fourier Feature Networks github.com
Official implementation of Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains.
-
-
-
-
-
-
- next page of items loading…
