flax

There are 4 repositories under flax topic.

transformers
huggingface / transformers
🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.
nlp natural-language-processing pytorch language-model tensorflow bert language-models pytorch-transformers nlp-library transformer model-hub pretrained-models jax flax seq2seq speech-recognition hacktoberfest python machine-learning deep-learning
Language:Python 124588
probml / pyprobml
Python code for "Probabilistic Machine learning" book by Kevin Murphy
blackjax colab flax jax jupyter-notebooks machine-learning numpyro pml probabilistic-programming pymc3 pyro pytorch tensorflow
Language:Jupyter Notebook 6242
FlaxEngine
FlaxEngine / FlaxEngine
Flax Engine – multi-platform 3D game engine
flax flaxengine engine game-engine gamedev multi-platform game-development c-sharp cpp flax-engine
Language:C# 5441
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.
chatbot deep-learning flax jax language-model large-language-models llama natural-language-processing transformer
Language:Python 2224
phlippe / uvadlc_notebooks
Repository of Jupyter notebook tutorials for teaching the Deep Learning Course at the University of Amsterdam (MSc AI), Fall 2023
deep-learning flax jax optax pytorch pytorch-lightning tutorial tutorials uvadlc
Language:Jupyter Notebook 2115
yang-song / score_sde
Official code for Score-Based Generative Modeling through Stochastic Differential Equations (ICLR 2021, Oral)
score-matching stochastic-differential-equations generative-models score-based-generative-modeling controllable-generation inverse-problems jax flax iclr-2021 diffusion-models
Language:Jupyter Notebook 1297
fortuna
awslabs / fortuna
A Library for Uncertainty Quantification.
ai bayesian-inference calibration conformal-prediction deep-learning flax jax machine-learning ml model-calibration neural-networks uncertainty uncertainty-calibration uncertainty-estimation uncertainty-quantification
Language:Python 849
google-research / long-range-arena
Long Range Arena for Benchmarking Efficient Transformers
transformers nlp deep-learning attention jax flax
Language:Python 680
google-research / prompt-tuning
Original Implementation of Prompt Tuning from Lester, et al, 2021
flax jax language-model machine-learning nlp prompt-tuning
Language:Python 619
ikostrikov / jaxrl
JAX (Flax) implementation of algorithms for Deep Reinforcement Learning with continuous action spaces.
batch-reinforcement-learning behavioral-cloning continuous-control deep-deterministic-policy-gradient deep-learning deep-reinforcement-learning flax gym jax offline-reinforcement-learning reinforcement-learning sac soft-actor-critic
Language:Jupyter Notebook 575
get-started-with-JAX
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.
deep-learning flax haiku jax jupyter lax learn-jax machine-learning numpy optax python tutorial xla
Language:Jupyter Notebook 557
texttron / tevatron
Tevatron - A flexible toolkit for neural retrieval research and development.
dense-retrieval dpr flax information-retrieval jax pytorch question-answering transformer
Language:Python 384
luyug / GradCache
Run Effective Large Batch Contrastive Learning Beyond GPU/TPU Memory Constraint
jax flax contrastive-learning pytorch deep-learning embedding large-language-models vector-database
Language:Python 304
google / CommonLoopUtils
CLU lets you write beautiful training loops in JAX.
flax jax ml
Language:Jupyter Notebook 293
nanodl
HMUNACHI / nanodl
A Jax-based library for designing and training transformer models from scratch.
attention-mechanism deep-learning distributed-training gpt jax llama machine-learning nlp transformer attention flax mistral
Language:Python 247
matthias-wright / flaxmodels
Pretrained deep learning models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet, etc.
stylegan2 gpt2 resnet vgg jax flax pretrained-models deep-learning
Language:Python 222
google / orbax
Orbax provides common utility libraries for JAX users.
checkpoint flax jax
Language:Python 221
Sea-Snell / JAXSeq
Train very large language models in Jax.
gpt2 gpt3 huggingface language-models opt deep-learning flax jax
Language:Python 181
EasyDeL
erfanzar / EasyDeL
EasyDeL is an OpenSource Library to make your training faster and more Optimized With cool Options for training and serving Both in Python And Mojo🔥
easydel flax gpt jax machine-learning mojo nlp optax pytorch transformers
Language:Python 147
flax-lang / flax
general purpose programming language, in the vein of C++
llvm language c-plus-plus language-design flax compiler programming-language
Language:C++ 145
DarshanDeshpande / jax-models
Unofficial JAX implementations of deep learning research papers
jax flax machine-learning artificial-intelligence research-paper-implementation transformers convolutional-neural-networks deep-learning computer-vision
Language:Python 139
FlaxAPI
FlaxEngine / FlaxAPI
Old repository with C# Editor and C# API for creating games in Flax Engine
game-engine gamedev flax flax-engine game-development api editor 3d-engine 3d-game-engine flax-editor games
Language:C# 134
jaketae / koclip
KoCLIP: Korean port of OpenAI CLIP, in Flax
flax jax openai-clip roberta vision-transformer
Language:Python 128
rwightman / efficientnet-jax
EfficientNet, MobileNetV3, MobileNetV2, MixNet, etc in JAX w/ Flax Linen and Objax
jax objax flax efficientnet mobilenetv3 mobilenetv2 mixnet tpu flax-linen
Language:Python 123
yang-song / score_flow
Official code for "Maximum Likelihood Training of Score-Based Diffusion Models", NeurIPS 2021 (spotlight)
score-based-generative-modeling score-based-generative-models diffusion-models score-matching jax flax normalizing-flows neurips-2021 neural-odes
Language:Python 118
Kaushalya / medclip
A multi-modal CLIP model trained on the medical dataset ROCO
transformers jax flax huggingface machine-learning medical-image-analysis
Language:Jupyter Notebook 106
n2cholas / jax-resnet
Implementations and checkpoints for ResNet, Wide ResNet, ResNeXt, ResNet-D, and ResNeSt in JAX (Flax).
computer-vision deep-learning flax jax machine-learning neural-networks resnet
Language:Python 100
evgenii-nikishin / rl_with_resets
JAX implementation of deep RL agents with resets from the paper "The Primacy Bias in Deep Reinforcement Learning"
atari deep-learning deep-reinforcement-learning dmc flax gym machine-learning overfitting reinforcement-learning sac spr atari2600 deepmind-control-suite drq jax soft-actor-critic
Language:Python 95
FlaxSamples
FlaxEngine / FlaxSamples
Collection of example projects for Flax Engine
flax flax-engine samples gamedev templates visual-studio
Language:C# 92
young-geng / m3ae_public
Multimodal Masked Autoencoders (M3AE): A JAX/Flax Implementation
computer-vision flax jax natural-language-processing transformers
Language:Python 91
shyamsn97 / hyper-nn
Easy Hypernetworks in Pytorch and Jax
flax hypernetworks jax machine-learning neural-networks pytorch
Language:Jupyter Notebook 90
FlaxDocs
FlaxEngine / FlaxDocs
Flax Engine Documentation
flax-documentation gamedev flax flax-engine game-engine documentation docs docfx flaxengine
Language:CSS 88
justinphan3110 / SciFive
SciFive: a text-text transformer model for biomedical literature
nlp transformer transfer-learning biomedical-language pubmed pmc attention pretrained-models jax flax huggingface
Language:Jupyter Notebook 85
google-research / jestimator
Amos optimizer with JEstimator lib.
deep-learning flax jax language-model machine-learning mnist optimization optimization-algorithms optimizer transformer
Language:Python 77
chef-transformer
chef-transformer / chef-transformer
Chef Transformer 🍲 .
recipes recipe-generator recipe-generation google-t5 t5 transformers huggingface jax flax
Language:Python 76
henry-prior / jax-rl
JAX implementations of core Deep RL algorithms
jax td3 reinforcement-learning deep-learning actor-critic deep-reinforcement-learning mujoco sac soft-actor-critic maximum-a-posteriori-estimation deepmind flax mpo
Language:Python 76