An open API service providing repository metadata for many open source software ecosystems.

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

Related Keywords
flax 244 jax 191 deep-learning 62 machine-learning 42 flax-engine 34 pytorch 24 python 23 transformers 19 transformer 18 reinforcement-learning 15 computer-vision 15 huggingface 15 optax 13 nlp 12 gamedev 11 tensorflow 10 llm 10 flaxengine 8 vision-transformer 8 neural-networks 8 tpu 7 language-model 7 haiku 6 deep-reinforcement-learning 6 diffusion-models 6 neural-network 6 sac 6 soft-actor-critic 5 tutorial 5 large-language-models 5 gpt 5 nnx 5 csharp 5 natural-language-processing 4 gym 4 score-based-generative-modeling 4 gpt2 4 kaggle 4 attention 4 diffusion 4 bert 3 game-engine 3 optimization 3 image-generation 3 colab 3 image-classification 3 td3 3 transfer-learning 3 numpyro 3 dm-haiku 3 ai 3 resnet 3 language-models 3 graph-neural-networks 3 gpt-2 3 generative-ai 3 numpy 3 offline-reinforcement-learning 3 pretrained-models 3 flax-linen 3 optimizer 3 gnn 3 wandb 2 gpu-acceleration 2 ddpm 2 digital-elevation-model 2 bayesian-inference 2 pytorch-lightning 2 cnn 2 distributed-training 2 attention-mechanism 2 nlp-library 2 clm 2 score-based-generative-models 2 score-matching 2 networking 2 information-retrieval 2 model-based-rl 2 model-based-reinforcement-learning 2 dqn 2 rl 2 ml 2 continuous-control 2 deep-deterministic-policy-gradient 2 unfinished 2 vae 2 segmentation 2 documentation 2 unet 2 docfx 2 neural-radiance-fields 2 recurrent-neural-networks 2 nerf 2 compter-vision 2 fine-tuning 2 flickr8k-dataset 2 tensor-networks 2 image-annotation 2 image-captioning 2 python3 2