GitHub topics: flax
google/orbax
Orbax provides common checkpointing and persistence utilities for JAX users
Language: Python - Size: 5.21 MB - Last synced at: about 1 hour ago - Pushed at: about 3 hours ago - Stars: 422 - Forks: 66

DBraun/jraphx
Graph Neural Networks with JAX/NNX
Language: Python - Size: 580 KB - Last synced at: 1 day ago - Pushed at: 1 day ago - Stars: 1 - Forks: 0

NeuralQXLab/nqxpack
Save/Load files from NetKet, flax and other scientific ML libraries
Language: Python - Size: 61.5 KB - Last synced at: 2 days ago - Pushed at: 3 days ago - Stars: 1 - Forks: 0

FlaxEngine/FlaxEngine
Flax Engine ā multi-platform 3D game engine
Language: C# - Size: 63.5 MB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 6,321 - Forks: 620

probml/pyprobml
Python code for "Probabilistic Machine learning" book by Kevin Murphy
Language: Jupyter Notebook - Size: 4.88 GB - Last synced at: about 23 hours ago - Pushed at: 9 months ago - Stars: 6,878 - Forks: 1,573

erfanzar/eformer
(EasyDel Former) is a utility library designed to simplify and enhance the development in JAX
Language: Python - Size: 3.86 MB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 28 - Forks: 4

bsc-quantic/tn4ml
Tensor Networks for Machine Learning
Language: Python - Size: 38.9 MB - Last synced at: 5 days ago - Pushed at: 5 days ago - Stars: 18 - Forks: 5

mashaan14/YouTube-channel
Posts related to my YouTube videos
Language: Jupyter Notebook - Size: 1.41 GB - Last synced at: 6 days ago - Pushed at: 6 days ago - Stars: 7 - Forks: 3

Locamage/jimm
JAX Image Modeling of Models contains Computer Vision/Vision Language Model implementations in native flax nnx with proper sharding annotations (allowing you to do easy fully sharded data parallel training with native flax nnx/jax)
Language: Python - Size: 7.76 MB - Last synced at: 6 days ago - Pushed at: 6 days ago - Stars: 3 - Forks: 0

FlaxEngine/FlaxDocs
Flax Engine Documentation
Language: CSS - Size: 762 MB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 105 - Forks: 69

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.
Language: Jupyter Notebook - Size: 1.78 MB - Last synced at: 7 days ago - Pushed at: almost 2 years ago - Stars: 748 - Forks: 113

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.
Language: Python - Size: 378 KB - Last synced at: 4 days ago - Pushed at: about 1 year ago - Stars: 2,497 - Forks: 261

mathisgerdes/bijx
Bijections & normalizing flows with JAX/NNX
Language: Python - Size: 1.91 MB - Last synced at: 12 days ago - Pushed at: 12 days ago - Stars: 0 - Forks: 0

ayushgabani/learn-jax
š Explore JAX for machine learning with clear notes on Just-In-Time compilation, automatic vectorization, and differentiation techniques.
Language: Jupyter Notebook - Size: 115 KB - Last synced at: 14 days ago - Pushed at: 14 days ago - Stars: 0 - Forks: 0

srigas/jaxKAN
Kolmogorov-Arnold Networks built on JAX
Language: Python - Size: 16.5 MB - Last synced at: 2 days ago - Pushed at: 5 months ago - Stars: 85 - Forks: 11

FlaxEngine/ArizonaFramework
Ready to use, open-source framework for creating games in Flax
Language: C++ - Size: 1.09 MB - Last synced at: 15 days ago - Pushed at: 3 months ago - Stars: 56 - Forks: 9

dirmeier/hierarchical-vqvae
A hierarchical VQ-VAE implementation in Flax
Language: Python - Size: 336 KB - Last synced at: 15 days ago - Pushed at: 15 days ago - Stars: 1 - Forks: 0

dirmeier/diffusion-transformer
A diffusion transformer implementation in Flax
Language: Python - Size: 210 KB - Last synced at: 15 days ago - Pushed at: 15 days ago - Stars: 2 - Forks: 0

dirmeier/block-neural-autoregressive-flow
A BNAF implementation in Flax
Language: Python - Size: 386 KB - Last synced at: 15 days ago - Pushed at: 15 days ago - Stars: 0 - Forks: 0

dirmeier/consistency-flow-matching
Flow and consistency matching in Flax
Language: Python - Size: 1.15 MB - Last synced at: 15 days ago - Pushed at: 15 days ago - Stars: 3 - Forks: 1

dirmeier/ramsey
Probabilistic deep learning using JAX
Language: Python - Size: 3.38 MB - Last synced at: 7 days ago - Pushed at: 7 months ago - Stars: 15 - Forks: 2

kousuke-nakano/orbkit
`orbkit` is a JAX-compatible toolkit for continuous ab initio quantum Monte Carlo (QMC) simulations, developed entirely from scratch using Python and JAX.
Language: Python - Size: 319 KB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 7 - Forks: 1

Auxeno/lotus
High-performance JAX (Flax) reinforcement learning library
Language: Python - Size: 1.56 MB - Last synced at: 18 days ago - Pushed at: 18 days ago - Stars: 2 - Forks: 0

dirmeier/blaxbird
A high-level API to build and train NNX models.
Language: Python - Size: 1.34 MB - Last synced at: 19 days ago - Pushed at: 19 days ago - Stars: 0 - Forks: 0

BobMcDear/flaim
Flax Image Models - State-of-the-art pre-trained vision backbones for Flax.
Language: Python - Size: 237 KB - Last synced at: 19 days ago - Pushed at: 3 months ago - Stars: 21 - Forks: 1

shyamsn97/hyper-nn
Easy Hypernetworks in Pytorch and Jax
Language: Jupyter Notebook - Size: 10.4 MB - Last synced at: 16 days ago - Pushed at: over 2 years ago - Stars: 104 - Forks: 6

Swoorup/rust-ecs-comparison
Comparison by example between different rust ecs libraries
Language: Rust - Size: 82 KB - Last synced at: 7 days ago - Pushed at: 21 days ago - Stars: 0 - Forks: 0

google/CommonLoopUtils
CLU lets you write beautiful training loops in JAX.
Language: Jupyter Notebook - Size: 1.49 MB - Last synced at: 3 days ago - Pushed at: 2 months ago - Stars: 355 - Forks: 33

epignatelli/helx
Interoperating between (Deep) Reiforcement Learning libraries
Language: Python - Size: 609 KB - Last synced at: 1 day ago - Pushed at: over 1 year ago - Stars: 9 - Forks: 6

SauravMaheshkar/mambax
mamba in JAX
Language: Python - Size: 13.7 KB - Last synced at: 8 days ago - Pushed at: 25 days ago - Stars: 0 - Forks: 0

texttron/tevatron
Tevatron - Unified Document Retrieval Toolkit across Scale, Language, and Modality. Demo in SIGIR 2023, SIGIR 2025.
Language: Python - Size: 20.5 MB - Last synced at: 27 days ago - Pushed at: 27 days ago - Stars: 678 - Forks: 113

MasterSkepticista/detr
JAX implementation of DETR
Language: Python - Size: 6.37 MB - Last synced at: 23 days ago - Pushed at: 2 months ago - Stars: 7 - Forks: 0

SauravMaheshkar/nanollm
JAX LLM playground
Language: Jupyter Notebook - Size: 85 KB - Last synced at: 7 days ago - Pushed at: about 1 month ago - Stars: 3 - Forks: 0

WICKED4950/BwETAFv2-130M
A new updated version of BwETA which uses JAX/Flax trained on TPU and has most of the advance features used by huge CLMs
Size: 3.91 KB - Last synced at: about 1 month ago - Pushed at: about 1 month ago - Stars: 0 - Forks: 0

WICKED4950/BwETAFv2-53M
A new updated version of BwETA which uses JAX/Flax trained on TPU and has most of the advance features used by huge CLMs
Size: 7.81 KB - Last synced at: about 1 month ago - Pushed at: about 1 month ago - Stars: 0 - Forks: 0

vietfood/banhxeo
A simple NLP library with educational purpose in mind.
Language: Python - Size: 6.36 MB - Last synced at: about 1 month ago - Pushed at: about 1 month ago - Stars: 30 - Forks: 3

HenryNdubuaku/nanodl
A Jax-based library for building transformers, includes implementations of GPT, Gemma, LlaMa, Mixtral, Whisper, SWin, ViT and more.
Language: Python - Size: 44.4 MB - Last synced at: 4 days ago - Pushed at: about 1 year ago - Stars: 290 - Forks: 11

alexander-koch/xmodality
Code for the paper "Cross-modality image synthesis from TOF-MRA to CTA using diffusion-based models"
Language: Python - Size: 3.97 MB - Last synced at: about 1 month ago - Pushed at: about 1 month ago - Stars: 4 - Forks: 1

matthias-wright/flaxmodels
Pretrained deep learning models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet, etc.
Language: Python - Size: 118 MB - Last synced at: about 1 month ago - Pushed at: 6 months ago - Stars: 255 - Forks: 25

ShawonAshraf/attention-visualiser
a module to visualise attention layer activations from transformer based models from huggingface
Language: Python - Size: 137 KB - Last synced at: 11 days ago - Pushed at: 2 months ago - Stars: 3 - Forks: 0

phlippe/uvadlc_notebooks
Repository of Jupyter notebook tutorials for teaching the Deep Learning Course at the University of Amsterdam (MSc AI), Fall 2023
Language: Jupyter Notebook - Size: 158 MB - Last synced at: about 2 months ago - Pushed at: 6 months ago - Stars: 2,919 - Forks: 633

google-research/jestimator š¦
Amos optimizer with JEstimator lib.
Language: Python - Size: 3.03 MB - Last synced at: 14 days ago - Pushed at: over 1 year ago - Stars: 82 - Forks: 6

BobMcDear/open-clip-jax
Implementation of OpenAI's CLIP in JAX/Flax
Language: Python - Size: 1.37 MB - Last synced at: about 2 months ago - Pushed at: about 2 months ago - Stars: 4 - Forks: 2

matthias-wright/jax-fid
FID computation in Jax/Flax.
Language: Python - Size: 55.9 MB - Last synced at: 21 days ago - Pushed at: about 1 year ago - Stars: 28 - Forks: 5

MineBill/Seed
A cross-platform launcher for the Flax game engine.
Language: C# - Size: 3.71 MB - Last synced at: 2 months ago - Pushed at: 2 months ago - Stars: 15 - Forks: 3

yonesuke/jaxfss
JAX/Flax implementation of finite-size scaling
Language: Python - Size: 1.93 MB - Last synced at: 13 days ago - Pushed at: over 2 years ago - Stars: 14 - Forks: 0

FlaxEngine/ArizonaFrameworkSample
Sample project showcasing usage of Arizona Framework for Flax Engine
Language: C# - Size: 2.44 MB - Last synced at: 2 months ago - Pushed at: 5 months ago - Stars: 22 - Forks: 4

ShawonAshraf/annotated-jax š¦
Deep Learning examples using the Jax ecosystem of libraries
Language: Jupyter Notebook - Size: 10.6 MB - Last synced at: 23 days ago - Pushed at: 4 months ago - Stars: 5 - Forks: 0

ziatdinovmax/NeuroBayes
Fully and Partially Bayesian Neural Nets
Language: Python - Size: 129 MB - Last synced at: 2 months ago - Pushed at: 4 months ago - Stars: 74 - Forks: 9

n2cholas/jax-resnet
Implementations and checkpoints for ResNet, Wide ResNet, ResNeXt, ResNet-D, and ResNeSt in JAX (Flax).
Language: Python - Size: 49.8 KB - Last synced at: 19 days ago - Pushed at: over 3 years ago - Stars: 112 - Forks: 8

FlaxEngine/FlaxSamples
Collection of example projects for Flax Engine
Language: C# - Size: 124 MB - Last synced at: 2 months ago - Pushed at: 2 months ago - Stars: 120 - Forks: 41

YVandana/NN-with-JAX
JAX Neural Network with Metrics and Visualization
Language: Python - Size: 7.81 KB - Last synced at: 2 months ago - Pushed at: 3 months ago - Stars: 0 - Forks: 0

BorealisAI/flora-opt
This is the official repository for the paper "Flora: Low-Rank Adapters Are Secretly Gradient Compressors" in ICML 2024.
Language: Python - Size: 672 KB - Last synced at: 4 days ago - Pushed at: about 1 year ago - Stars: 104 - Forks: 5

yang-song/score_flow
Official code for "Maximum Likelihood Training of Score-Based Diffusion Models", NeurIPS 2021 (spotlight)
Language: Python - Size: 65.4 KB - Last synced at: about 2 months ago - Pushed at: over 3 years ago - Stars: 144 - Forks: 24

ikostrikov/jaxrl
JAX (Flax) implementation of algorithms for Deep Reinforcement Learning with continuous action spaces.
Language: Jupyter Notebook - Size: 4.17 MB - Last synced at: 3 months ago - Pushed at: almost 3 years ago - Stars: 685 - Forks: 72

FlaxEngine/Bistro
Amazon Lumberyard Bistro for Flax Engine
Language: C# - Size: 781 MB - Last synced at: 3 months ago - Pushed at: 3 months ago - Stars: 5 - Forks: 0

Kaushalya/medclip
A multi-modal CLIP model trained on the medical dataset ROCO
Language: Jupyter Notebook - Size: 419 MB - Last synced at: 3 months ago - Pushed at: 3 months ago - Stars: 138 - Forks: 19

mathpluscode/ImgX-DiffSeg
A JAX-based deep learning framework for image segmentation using diffusion models.
Language: Python - Size: 26.4 MB - Last synced at: 3 months ago - Pushed at: about 1 year ago - Stars: 84 - Forks: 8

camenduru/stable-diffusion-diffusers-colab
š¤ HuggingFace Diffusers Flax TPU and PyTorch GPU for Colab
Language: Jupyter Notebook - Size: 121 KB - Last synced at: 7 days ago - Pushed at: almost 2 years ago - Stars: 39 - Forks: 6

jeongsol-kim/jax_practice
Jax/Flax implementation practice for various neural networks
Language: Python - Size: 23.4 MB - Last synced at: about 1 month ago - Pushed at: about 3 years ago - Stars: 4 - Forks: 0

DBraun/librosax
Librosa in JAX
Language: Python - Size: 259 KB - Last synced at: about 18 hours ago - Pushed at: 17 days ago - Stars: 2 - Forks: 0

JoeJoe1313/PaliGemma-Image-Segmentation
An app with FastAPI, Docker, transformers, JAX/Flax for performing image segmentation with PaliGemma 2 mix
Language: Jupyter Notebook - Size: 8.78 MB - Last synced at: 3 months ago - Pushed at: 3 months ago - Stars: 0 - Forks: 0

erfanzar/FJDiffusion
implementation of in Jax/Flax !
Language: Python - Size: 136 KB - Last synced at: about 1 month ago - Pushed at: about 2 years ago - Stars: 5 - Forks: 0

jaco-bro/diajax
Dia-JAX: A JAX port of Dia, the text-to-speech model for generating realistic dialogue from text with emotion and tone control.
Language: Jupyter Notebook - Size: 914 KB - Last synced at: about 2 months ago - Pushed at: 4 months ago - Stars: 27 - Forks: 4

ejlnmusic/PaliGemma-flickr8k-finetuning
# PaliGemma-flickr8k-finetuningThis repository provides a method to fine-tune the PaliGemma model on the Flickr8k dataset for improved image captioning. Explore the features and utilities designed for efficient training and testing. šš
Language: Jupyter Notebook - Size: 375 KB - Last synced at: 3 months ago - Pushed at: 3 months ago - Stars: 0 - Forks: 0

AHMEDSANA/PaliGemma-flickr8k-finetuning
This repository contains code for fine-tuning Google's PaliGemma vision-language model on the Flickr8k dataset for image captioning tasks
Language: Jupyter Notebook - Size: 401 KB - Last synced at: 3 months ago - Pushed at: 3 months ago - Stars: 1 - Forks: 0

yang-song/score_sde
Official code for Score-Based Generative Modeling through Stochastic Differential Equations (ICLR 2021, Oral)
Language: Jupyter Notebook - Size: 4.35 MB - Last synced at: 4 months ago - Pushed at: almost 3 years ago - Stars: 1,659 - Forks: 222

ghubnerr/attention-mechanisms
A compilation of most State-of-the-Art Attention Mechanisms: MHSA, MQA, GQA, and MLA, implemented in JAX
Language: Python - Size: 129 KB - Last synced at: 2 months ago - Pushed at: 4 months ago - Stars: 2 - Forks: 0

google-research/prompt-tuning
Original Implementation of Prompt Tuning from Lester, et al, 2021
Language: Python - Size: 3.41 MB - Last synced at: 4 months ago - Pushed at: 6 months ago - Stars: 680 - Forks: 60

XanaduAI/GradDFT
GradDFT is a JAX-based library enabling the differentiable design and experimentation of exchange-correlation functionals using machine learning techniques.
Language: Python - Size: 181 MB - Last synced at: about 2 months ago - Pushed at: over 1 year ago - Stars: 101 - Forks: 9

jaketae/koclip
KoCLIP: Korean port of OpenAI CLIP, in Flax
Language: Python - Size: 27.9 MB - Last synced at: 4 months ago - Pushed at: about 2 years ago - Stars: 151 - Forks: 18

xmarva/jax-vit
ViT adaptation with Jax/Flax
Language: Jupyter Notebook - Size: 122 KB - Last synced at: 4 months ago - Pushed at: 4 months ago - Stars: 0 - Forks: 0

nobodyPerfecZ/vit-finetuned-patch16-224-recaptchav2-v1
An AI-powered computer vision system written in PyTorch, Tensorflow and JAX, designed to analyze and solve reCAPTCHA v2 challenges.
Language: Jupyter Notebook - Size: 7.81 KB - Last synced at: 4 months ago - Pushed at: 4 months ago - Stars: 0 - Forks: 0

nipunsadvilkar/roberta-base-mr
RoBERTa Marathi Language model trained from scratch during huggingface š¤ x flax community week
Language: Python - Size: 440 KB - Last synced at: 19 days ago - Pushed at: about 4 years ago - Stars: 28 - Forks: 4

luyug/GradCache
Run Effective Large Batch Contrastive Learning Beyond GPU/TPU Memory Constraint
Language: Python - Size: 43.9 KB - Last synced at: 4 months ago - Pushed at: over 1 year ago - Stars: 387 - Forks: 24

flax-lang/flax
general purpose programming language, in the vein of C++
Language: C++ - Size: 5.25 MB - Last synced at: 4 months ago - Pushed at: about 4 years ago - Stars: 149 - Forks: 10

AshishKumar4/FlaxDiff
A simple, easy-to-understand library for diffusion models using Flax and Jax. Includes detailed notebooks on DDPM, DDIM, and EDM with simplified mathematical explanations. Made as part of my journey for learning and experimenting with generative AI.
Language: Jupyter Notebook - Size: 238 MB - Last synced at: 4 months ago - Pushed at: 4 months ago - Stars: 25 - Forks: 0

frankroeder/goal_conditioned_rl
Goal-conditioned reinforcement learning like š„
Language: Python - Size: 29.3 KB - Last synced at: 5 days ago - Pushed at: over 1 year ago - Stars: 12 - Forks: 0

EmptyJackson/unifloral
Unified Implementations of Offline Reinforcement Learning Algorithms
Language: Python - Size: 47.9 KB - Last synced at: 4 months ago - Pushed at: 4 months ago - Stars: 58 - Forks: 3

JeyRunner/flaxfit
Fitting jax flax models made simple.
Language: Python - Size: 138 KB - Last synced at: 4 months ago - Pushed at: 4 months ago - Stars: 0 - Forks: 0

DarshanDeshpande/jax-models
Unofficial JAX implementations of deep learning research papers
Language: Python - Size: 201 KB - Last synced at: about 2 months ago - Pushed at: about 3 years ago - Stars: 156 - Forks: 9

awslabs/fortuna š¦
A Library for Uncertainty Quantification.
Language: Python - Size: 4.56 MB - Last synced at: 5 months ago - Pushed at: 5 months ago - Stars: 913 - Forks: 48

FlaxEngine/NetworkSample
Multiplayer sample with players lobby and chat in Flax Engine
Language: C# - Size: 1.86 MB - Last synced at: 5 days ago - Pushed at: 5 months ago - Stars: 21 - Forks: 7

Crawcik/Journal
Console for Flax Engine with command handling
Language: C# - Size: 117 KB - Last synced at: 5 months ago - Pushed at: 5 months ago - Stars: 8 - Forks: 0

Sea-Snell/JAXSeq
Train very large language models in Jax.
Language: Python - Size: 252 KB - Last synced at: 4 months ago - Pushed at: almost 2 years ago - Stars: 204 - Forks: 18

FlaxEngine/OnlinePlatformXboxLive
Online platform implementation of Xbox Live for games made with Flax Engine
Language: C++ - Size: 19.5 KB - Last synced at: 5 months ago - Pushed at: 5 months ago - Stars: 4 - Forks: 1

evanatyourservice/image-classification-jax
Image classification in JAX with ViT, resnet, cifar10, cifar100, imagenette, and imagenet
Language: Python - Size: 64.5 KB - Last synced at: 14 days ago - Pushed at: 8 months ago - Stars: 3 - Forks: 0

NobuoTsukamoto/jax_examples
Jax, Flax, examples (ImageClassification, SemanticSegmentation, and more...)
Language: Python - Size: 4.35 MB - Last synced at: 4 months ago - Pushed at: 5 months ago - Stars: 10 - Forks: 0

evgenii-nikishin/rl_with_resets
JAX implementation of deep RL agents with resets from the paper "The Primacy Bias in Deep Reinforcement Learning"
Language: Python - Size: 1.1 MB - Last synced at: 24 days ago - Pushed at: over 3 years ago - Stars: 100 - Forks: 7

ghubnerr/machine-learning
A dump to all things LLMs, CUDA, DL frameworks, Tokenization Methods, and latest stuff Iām building/researching with
Language: Jupyter Notebook - Size: 8.64 MB - Last synced at: 3 months ago - Pushed at: 5 months ago - Stars: 0 - Forks: 0

FlaxEngine/ImGui
Dear ImGui plugin for Flax Engine that adds debug GUI interface to game viewport.
Language: C++ - Size: 1.16 MB - Last synced at: 3 months ago - Pushed at: 5 months ago - Stars: 13 - Forks: 5

terrence-ou/Generative-Deep-Learning-2nd-Edition-PyTorch-JAX š¦
Pytorch/JAX implementation of examples in Generative Deep Learning 2nd Edition by David Foster
Language: Jupyter Notebook - Size: 31.2 MB - Last synced at: 5 months ago - Pushed at: almost 2 years ago - Stars: 21 - Forks: 4

sooheon/perceiver-jax
Perceiver (transformer variant) implemented in JAX and Flax
Language: Python - Size: 10.7 KB - Last synced at: about 2 months ago - Pushed at: over 4 years ago - Stars: 12 - Forks: 2

NITHISHM2410/flax-pilot
A Simplistic trainer for Flax
Language: Python - Size: 1.62 MB - Last synced at: 5 months ago - Pushed at: 5 months ago - Stars: 7 - Forks: 0

satojkovic/gpt-tf-pytorch-jax
GPT from scratch (tensorflow / pytorch / jax)
Language: Python - Size: 28.3 KB - Last synced at: 5 months ago - Pushed at: 5 months ago - Stars: 0 - Forks: 0

alexjackson1/jax-transformer
A re-implementation of Callum McDougall's "Transformer from Scratch" using JAX (Flax).
Language: Python - Size: 52.7 KB - Last synced at: 5 months ago - Pushed at: 5 months ago - Stars: 0 - Forks: 0

young-geng/m3ae_public
Multimodal Masked Autoencoders (M3AE): A JAX/Flax Implementation
Language: Python - Size: 28.3 KB - Last synced at: 3 months ago - Pushed at: 6 months ago - Stars: 103 - Forks: 12

IvanIsCoding/GNN-for-Combinatorial-Optimization
JAX + Flax implementation of "Combinatorial Optimization with Physics-Inspired Graph Neural Networks" by Schuetz et al.
Language: Jupyter Notebook - Size: 3.31 MB - Last synced at: 5 months ago - Pushed at: over 2 years ago - Stars: 50 - Forks: 4

dtunai/Griffin-Jax
Jax implementation of "Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models"
Language: Python - Size: 53.7 KB - Last synced at: 2 months ago - Pushed at: over 1 year ago - Stars: 14 - Forks: 0

rwightman/efficientnet-jax
EfficientNet, MobileNetV3, MobileNetV2, MixNet, etc in JAX w/ Flax Linen and Objax
Language: Python - Size: 157 KB - Last synced at: about 2 months ago - Pushed at: over 1 year ago - Stars: 128 - Forks: 14
