flax

Python code for "Machine learning: a probabilistic perspective" (2nd edition)

Flax Engine – multi-platform 3D game engine

Official code for Score-Based Generative Modeling through Stochastic Differential Equations (ICLR 2021, Oral)

Long Range Arena for Benchmarking Efficient Transformers

Jax (Flax) implementation of algorithms for Deep Reinforcement Learning with continuous action spaces.

general purpose programming language, in the vein of C++

Old repository with C# Editor and C# API for creating games in Flax Engine

EfficientNet, MobileNetV3, MobileNetV2, MixNet, etc in JAX w/ Flax Linen and Objax

Flax Engine Documentation

KoCLIP: Korean port of OpenAI CLIP, in Flax

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

CLU lets you write beautiful training loops in JAX.

Collection of example projects for Flax Engine

Implementations and checkpoints for ResNet, Wide ResNet, ResNeXt, ResNet-D, and ResNeSt in JAX (Flax).

JAX implementations of core Deep RL algorithms

JAX code for the paper "Control-Oriented Model-Based Reinforcement Learning with Implicit Differentiation"

RoBERTa Marathi Language model trained from scratch during huggingface 🤗 x flax community week

A collection of optimizers, some arcane others well known, for Flax.

Jax implementation of the AdaHessian optimizer

Chef Transformer 🍲 .

Unofficial implementation of NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis, using Flax with the Linen API

A pathway and collection of resources to learning Jax from beginning to advance.

Flax (JAX) implementation of Progressive Growing of GANs for Improved Quality, Stability, and Variation

Perceiver (transformer variant) implemented in JAX and Flax

Implementation of Vision Transformers in Flax

Model Agnostic Meta Learning (MAML) implemented in Flax, the neural network library for JAX.

Basic VR Module for Flax

Implementation of several Generative Adversarial Networks in JAX / Flax