flax
There are 226 repositories under flax topic.
huggingface/transformers
🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.
huggingface/diffusers
🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.
probml/pyprobml
Python code for "Probabilistic Machine learning" book by Kevin Murphy
FlaxEngine/FlaxEngine
Flax Engine – multi-platform 3D game engine
phlippe/uvadlc_notebooks
Repository of Jupyter notebook tutorials for teaching the Deep Learning Course at the University of Amsterdam (MSc AI), Fall 2023
young-geng/EasyLM
Large language models (LLMs) made easy, EasyLM is a one stop solution for pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.
yang-song/score_sde
Official code for Score-Based Generative Modeling through Stochastic Differential Equations (ICLR 2021, Oral)
jla524/fromthetensor
From the Tensor to Stable Diffusion, a rough outline for a 1 week course.
awslabs/fortuna
A Library for Uncertainty Quantification.
google-research/long-range-arena
Long Range Arena for Benchmarking Efficient Transformers
gordicaleksa/get-started-with-JAX
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.
ikostrikov/jaxrl
JAX (Flax) implementation of algorithms for Deep Reinforcement Learning with continuous action spaces.
google-research/prompt-tuning
Original Implementation of Prompt Tuning from Lester, et al, 2021
texttron/tevatron
Tevatron - Unified Document Retrieval Toolkit across Scale, Language, and Modality. Demo in SIGIR 2023, SIGIR 2025.
google/orbax
Orbax provides common checkpointing and persistence utilities for JAX users
luyug/GradCache
Run Effective Large Batch Contrastive Learning Beyond GPU/TPU Memory Constraint
google/CommonLoopUtils
CLU lets you write beautiful training loops in JAX.
erfanzar/EasyDeL
Accelerate, Optimize performance with streamlined training and serving options with JAX.
HMUNACHI/nanodl
A Jax-based library for designing and training transformer models from scratch.
matthias-wright/flaxmodels
Pretrained deep learning models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet, etc.
Sea-Snell/JAXSeq
Train very large language models in Jax.
DarshanDeshpande/jax-models
Unofficial JAX implementations of deep learning research papers
nissymori/JAX-CORL
Clean single-file implementation of offline RL algorithms in JAX
yang-song/score_flow
Official code for "Maximum Likelihood Training of Score-Based Diffusion Models", NeurIPS 2021 (spotlight)
jaketae/koclip
KoCLIP: Korean port of OpenAI CLIP, in Flax
flax-lang/flax
general purpose programming language, in the vein of C++
Kaushalya/medclip
A multi-modal CLIP model trained on the medical dataset ROCO
FlaxEngine/FlaxAPI
Old repository with C# Editor and C# API for creating games in Flax Engine
rwightman/efficientnet-jax
EfficientNet, MobileNetV3, MobileNetV2, MixNet, etc in JAX w/ Flax Linen and Objax
FlaxEngine/FlaxSamples
Collection of example projects for Flax Engine
n2cholas/jax-resnet
Implementations and checkpoints for ResNet, Wide ResNet, ResNeXt, ResNet-D, and ResNeSt in JAX (Flax).
FlaxEngine/FlaxDocs
Flax Engine Documentation
BorealisAI/flora-opt
This is the official repository for the paper "Flora: Low-Rank Adapters Are Secretly Gradient Compressors" in ICML 2024.
shyamsn97/hyper-nn
Easy Hypernetworks in Pytorch and Jax
young-geng/m3ae_public
Multimodal Masked Autoencoders (M3AE): A JAX/Flax Implementation
evgenii-nikishin/rl_with_resets
JAX implementation of deep RL agents with resets from the paper "The Primacy Bias in Deep Reinforcement Learning"